chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.activation import (
|
||||
FastGELU,
|
||||
FatreluAndMul,
|
||||
GeluAndMul,
|
||||
MulAndSilu,
|
||||
NewGELU,
|
||||
QuickGELU,
|
||||
ReLUSquaredActivation,
|
||||
SiluAndMul,
|
||||
SiluAndMulWithClamp,
|
||||
SwigluOAIAndMul,
|
||||
SwigluStepAndMul,
|
||||
swiglustep_and_mul_triton,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16, torch.float]
|
||||
NUM_TOKENS = [7, 83, 2048] # Arbitrary values for testing
|
||||
D = [512, 13824] # Arbitrary values for testing
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"activation",
|
||||
[
|
||||
"silu_and_mul",
|
||||
"mul_and_silu",
|
||||
"gelu",
|
||||
"gelu_tanh",
|
||||
"fatrelu",
|
||||
"swigluoai_and_mul",
|
||||
"swiglustep_and_mul",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_act_and_mul(
|
||||
default_vllm_config,
|
||||
activation: str,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
|
||||
if activation == "silu_and_mul":
|
||||
layer = SiluAndMul(compile_native=False)
|
||||
fn = torch.ops._C.silu_and_mul
|
||||
if activation == "mul_and_silu":
|
||||
layer = MulAndSilu()
|
||||
fn = torch.ops._C.mul_and_silu
|
||||
elif activation == "gelu":
|
||||
layer = GeluAndMul(approximate="none")
|
||||
fn = torch.ops._C.gelu_and_mul
|
||||
elif activation == "gelu_tanh":
|
||||
layer = GeluAndMul(approximate="tanh")
|
||||
fn = torch.ops._C.gelu_tanh_and_mul
|
||||
elif activation == "fatrelu":
|
||||
threshold = random.uniform(0, 1)
|
||||
layer = FatreluAndMul(threshold)
|
||||
fn = torch.ops._C.fatrelu_and_mul
|
||||
elif activation == "swigluoai_and_mul":
|
||||
layer = SwigluOAIAndMul()
|
||||
fn = torch.ops._C.swigluoai_and_mul
|
||||
elif activation == "swiglustep_and_mul":
|
||||
layer = SwigluStepAndMul()
|
||||
fn = swiglustep_and_mul_triton
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]:
|
||||
rtol = {
|
||||
# For fp16, change the relative tolerance from 1e-3 to 2e-3
|
||||
torch.float16: 2e-3,
|
||||
torch.bfloat16: 2e-2,
|
||||
torch.float: 1.3e-6,
|
||||
}
|
||||
|
||||
def _get_rtol(output) -> float:
|
||||
return rtol[output.dtype]
|
||||
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=_get_rtol(out)
|
||||
)
|
||||
else:
|
||||
# The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are
|
||||
# equivalent to the native PyTorch implementations, so we can do exact
|
||||
# comparison.
|
||||
torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0)
|
||||
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = x.shape[:-1] + (d,)
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
if activation == "fatrelu":
|
||||
opcheck(fn, (out, x, threshold))
|
||||
elif activation == "swigluoai_and_mul":
|
||||
opcheck(fn, (out, x, layer.alpha, layer.limit))
|
||||
elif activation != "swiglustep_and_mul":
|
||||
opcheck(fn, (out, x))
|
||||
|
||||
|
||||
SWIGLU_LIMITS = [3.0, 7.0, 15.0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_and_mul_with_clamp(
|
||||
default_vllm_config,
|
||||
swiglu_limit: float,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
"""SiluAndMulWithClamp: cuda kernel must match native reference."""
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
# Use large values to ensure clamping is exercised.
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
|
||||
|
||||
layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
|
||||
rtol = {
|
||||
torch.float16: 2e-3,
|
||||
torch.bfloat16: 2e-2,
|
||||
torch.float: 1.3e-6,
|
||||
}
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
|
||||
)
|
||||
|
||||
# Verify clamping is actually being applied: the clamped output should
|
||||
# differ from the unclamped SiluAndMul output when inputs are large.
|
||||
unclamped_out = SiluAndMul.forward_native(x)
|
||||
assert not torch.equal(ref_out.float(), unclamped_out.float()), (
|
||||
"Input was not large enough to exercise the clamp; increase scale"
|
||||
)
|
||||
|
||||
# Verify gate clamping semantics with a controlled scalar case.
|
||||
# gate=large_val is clamped to limit first, then silu(limit) * 1.0.
|
||||
x_gate = torch.tensor(
|
||||
[[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
|
||||
expected_gate = torch.nn.functional.silu(
|
||||
torch.tensor(swiglu_limit, dtype=torch.float32)
|
||||
).item()
|
||||
torch.testing.assert_close(
|
||||
out_gate,
|
||||
torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# Verify up clamping semantics: up >> limit gets clamped to limit.
|
||||
x_up = torch.tensor(
|
||||
[[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
|
||||
silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
|
||||
torch.testing.assert_close(
|
||||
out_up,
|
||||
torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# opcheck
|
||||
out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
|
||||
opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"activation",
|
||||
[
|
||||
(FastGELU, torch.ops._C.gelu_fast),
|
||||
(NewGELU, torch.ops._C.gelu_new),
|
||||
(QuickGELU, torch.ops._C.gelu_quick),
|
||||
(ReLUSquaredActivation, torch.ops._C.relu_squared),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_activation(
|
||||
default_vllm_config,
|
||||
activation: type[torch.nn.Module],
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
x = torch.randn(num_tokens, d, dtype=dtype)
|
||||
layer = activation[0]()
|
||||
fn = activation[1]
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
|
||||
out = torch.empty_like(x)
|
||||
opcheck(fn, (out, x))
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for ApplyRotaryEmb CustomOp dispatch behavior.
|
||||
|
||||
This test ensures that RotaryEmbedding classes correctly call the appropriate
|
||||
ApplyRotaryEmb methods based on the calling context:
|
||||
|
||||
1. RotaryEmbedding.forward_native() -> ApplyRotaryEmb.forward_native()
|
||||
2. RotaryEmbedding.forward_cuda() -> ApplyRotaryEmb.forward() (auto-dispatch)
|
||||
3. RotaryEmbedding.forward_hip() -> ApplyRotaryEmb.forward() (auto-dispatch)
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
VllmConfig,
|
||||
get_cached_compilation_config,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RotaryEmbeddingTestCase:
|
||||
"""Test case configuration for RotaryEmbedding dispatch tests."""
|
||||
|
||||
name: str
|
||||
rope_class: type
|
||||
rope_kwargs: dict
|
||||
method_name: str # forward_native, forward_cuda, forward
|
||||
positions_shape: tuple # (num_tokens,) or (3, num_tokens) or (4, num_tokens)
|
||||
expect_forward_native: bool # Should call ApplyRotaryEmb.forward_native()
|
||||
expect_forward: bool # Should call ApplyRotaryEmb.forward()
|
||||
|
||||
|
||||
def get_test_cases() -> list[RotaryEmbeddingTestCase]:
|
||||
"""Generate test cases for all RotaryEmbedding classes."""
|
||||
from vllm.model_executor.layers.rotary_embedding.ernie45_vl_rope import (
|
||||
Ernie4_5_VLRotaryEmbedding,
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
|
||||
from vllm.model_executor.layers.rotary_embedding.xdrope import XDRotaryEmbedding
|
||||
|
||||
common_kwargs = {
|
||||
"head_size": 128,
|
||||
"rotary_dim": 128,
|
||||
"max_position_embeddings": 4096,
|
||||
"base": 10000,
|
||||
"is_neox_style": True,
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
|
||||
return [
|
||||
# MRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="MRotaryEmbedding.forward_native",
|
||||
rope_class=MRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
|
||||
method_name="forward_native",
|
||||
positions_shape=(3, 32), # 2D for multimodal
|
||||
expect_forward_native=True,
|
||||
expect_forward=False,
|
||||
),
|
||||
RotaryEmbeddingTestCase(
|
||||
name="MRotaryEmbedding.forward_cuda_1d",
|
||||
rope_class=MRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [16, 24, 24]},
|
||||
method_name="forward_cuda",
|
||||
positions_shape=(32,), # 1D triggers apply_rotary_emb path
|
||||
expect_forward_native=False,
|
||||
expect_forward=True,
|
||||
),
|
||||
# XDRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="XDRotaryEmbedding.forward",
|
||||
rope_class=XDRotaryEmbedding,
|
||||
rope_kwargs={
|
||||
**common_kwargs,
|
||||
"scaling_alpha": 1.0,
|
||||
"xdrope_section": [16, 16, 16, 16],
|
||||
},
|
||||
method_name="forward",
|
||||
positions_shape=(4, 32), # 4D for P/W/H/T
|
||||
expect_forward_native=False,
|
||||
expect_forward=True,
|
||||
),
|
||||
# Ernie4_5_VLRotaryEmbedding tests
|
||||
RotaryEmbeddingTestCase(
|
||||
name="Ernie4_5_VLRotaryEmbedding.forward_native",
|
||||
rope_class=Ernie4_5_VLRotaryEmbedding,
|
||||
rope_kwargs={**common_kwargs, "mrope_section": [22, 22, 20]},
|
||||
method_name="forward_native",
|
||||
positions_shape=(3, 32), # 2D for multimodal
|
||||
expect_forward_native=True,
|
||||
expect_forward=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def run_dispatch_test(
|
||||
test_case: RotaryEmbeddingTestCase,
|
||||
device: str,
|
||||
):
|
||||
"""Run a dispatch test for a RotaryEmbedding class."""
|
||||
vllm_config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"])
|
||||
)
|
||||
get_cached_compilation_config.cache_clear()
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
rope = test_case.rope_class(**test_case.rope_kwargs).to(device=device)
|
||||
|
||||
apply_rotary_emb = rope.apply_rotary_emb
|
||||
|
||||
# Verify custom op is enabled
|
||||
if test_case.expect_forward_native:
|
||||
assert (
|
||||
apply_rotary_emb._forward_method != apply_rotary_emb.forward_native
|
||||
), "Test setup error: ApplyRotaryEmb custom op should be enabled"
|
||||
|
||||
# Setup call tracking
|
||||
call_tracker = {"forward_native_called": False, "forward_called": False}
|
||||
original_forward_native = apply_rotary_emb.forward_native
|
||||
original_forward = apply_rotary_emb.forward
|
||||
|
||||
def tracked_forward_native(*args, **kwargs):
|
||||
call_tracker["forward_native_called"] = True
|
||||
return original_forward_native(*args, **kwargs)
|
||||
|
||||
def tracked_forward(*args, **kwargs):
|
||||
call_tracker["forward_called"] = True
|
||||
return original_forward(*args, **kwargs)
|
||||
|
||||
apply_rotary_emb.forward_native = tracked_forward_native
|
||||
apply_rotary_emb.forward = tracked_forward
|
||||
|
||||
try:
|
||||
num_tokens = test_case.positions_shape[-1]
|
||||
num_q_heads = 8
|
||||
num_kv_heads = 2
|
||||
head_size = test_case.rope_kwargs["head_size"]
|
||||
max_position = test_case.rope_kwargs["max_position_embeddings"]
|
||||
|
||||
positions = torch.randint(
|
||||
0, max_position // 4, test_case.positions_shape, device=device
|
||||
)
|
||||
query = torch.randn(
|
||||
num_tokens, num_q_heads * head_size, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
key = torch.randn(
|
||||
num_tokens,
|
||||
num_kv_heads * head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Call the method under test
|
||||
method = getattr(rope, test_case.method_name)
|
||||
method(positions, query.clone(), key.clone())
|
||||
|
||||
# Verify expectations
|
||||
if test_case.expect_forward_native:
|
||||
assert call_tracker["forward_native_called"], (
|
||||
f"{test_case.name} should call ApplyRotaryEmb.forward_native()"
|
||||
)
|
||||
if not test_case.expect_forward:
|
||||
assert not call_tracker["forward_called"], (
|
||||
f"{test_case.name} should NOT call ApplyRotaryEmb.forward(). "
|
||||
"Bug: when +apply_rotary_emb is enabled, forward_native() "
|
||||
"incorrectly dispatches to CUDA/HIP kernels."
|
||||
)
|
||||
if test_case.expect_forward:
|
||||
assert call_tracker["forward_called"], (
|
||||
f"{test_case.name} should call ApplyRotaryEmb.forward()"
|
||||
)
|
||||
finally:
|
||||
apply_rotary_emb.forward_native = original_forward_native
|
||||
apply_rotary_emb.forward = original_forward
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda tc: tc.name)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_rotary_embedding_dispatch(
|
||||
test_case: RotaryEmbeddingTestCase,
|
||||
device: str,
|
||||
):
|
||||
"""
|
||||
Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method.
|
||||
|
||||
- forward_native methods should call ApplyRotaryEmb.forward_native()
|
||||
- forward_cuda/forward methods should call ApplyRotaryEmb.forward()
|
||||
"""
|
||||
run_dispatch_test(test_case, device)
|
||||
@@ -0,0 +1,70 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the batched-weight RMS norm kernel (vllm._custom_ops.rms_norm).
|
||||
|
||||
``rms_norm`` can use the outermost input batch index to select the corresponding
|
||||
weight row. The result must match that of looping ``rms_norm`` over that dimension.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="rms_norm requires a CUDA/ROCm device",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(28, 17, 128), # 3D: [num_rows, tokens, hidden]
|
||||
(1, 5, 2, 128), # 4D: single row (edge case)
|
||||
(28, 13, 8, 128), # 4D: [L, num_ctx, nkv, hd] (DFlash K-norm)
|
||||
(6, 3, 4, 769), # 4D: non-power-of-two hidden size
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_matches_loop(
|
||||
shape: tuple[int, ...], dtype: torch.dtype, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
num_rows, hidden = shape[0], shape[-1]
|
||||
eps = 1e-6
|
||||
|
||||
x = torch.randn(*shape, dtype=dtype) * 0.1
|
||||
# Distinct weight per row so that a wrong row index would be caught.
|
||||
weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0
|
||||
|
||||
# Reference batched-weight rms norm.
|
||||
out_ref = torch.empty_like(x)
|
||||
for i in range(x.shape[0]):
|
||||
ops.rms_norm(out_ref[i], x[i], weight[i], eps)
|
||||
|
||||
out = torch.empty_like(x)
|
||||
ops.rms_norm(out, x, weight, eps)
|
||||
|
||||
# Expect bitwise-identical results.
|
||||
torch.testing.assert_close(out, out_ref, atol=0, rtol=0)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_validates_shapes() -> None:
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
x = torch.randn(4, 8, 128, dtype=torch.float)
|
||||
out = torch.empty_like(x)
|
||||
# Expect num rows mismatch.
|
||||
with pytest.raises(RuntimeError):
|
||||
ops.rms_norm(out, x, torch.randn(3, 128), 1e-6)
|
||||
# Expect hidden size mismatch.
|
||||
with pytest.raises(RuntimeError):
|
||||
ops.rms_norm(out, x, torch.randn(4, 64), 1e-6)
|
||||
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
from vllm.model_executor.layers.activation import (
|
||||
GELU,
|
||||
FastGELU,
|
||||
GeluAndMul,
|
||||
NewGELU,
|
||||
QuickGELU,
|
||||
SiluAndMul,
|
||||
)
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float32]
|
||||
NUM_TOKENS = [7, 83]
|
||||
D = [512, 2048]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("activation_cls", "fn"),
|
||||
[
|
||||
(SiluAndMul, torch.ops._C.silu_and_mul),
|
||||
(GeluAndMul, torch.ops._C.gelu_and_mul),
|
||||
(GeluAndMul, torch.ops._C.gelu_tanh_and_mul),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_act_and_mul(
|
||||
default_vllm_config,
|
||||
activation_cls: type[torch.nn.Module],
|
||||
fn: object,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
|
||||
|
||||
layer = activation_cls()
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
|
||||
output_shape = x.shape[:-1] + (x.shape[-1] // 2,)
|
||||
raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
opcheck(fn, (raw_out, x))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("activation_cls", "fn", "op_args"),
|
||||
[
|
||||
(NewGELU, torch.ops._C.gelu_new, ()),
|
||||
(FastGELU, torch.ops._C.gelu_fast, ()),
|
||||
(QuickGELU, torch.ops._C.gelu_quick, ()),
|
||||
pytest.param(
|
||||
GELU,
|
||||
getattr(torch.ops._C, "activation_lut_bf16", None),
|
||||
("gelu",),
|
||||
marks=pytest.mark.skipif(
|
||||
current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
|
||||
reason="activation_lut_bf16 is only built on Arm CPU",
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_unary_activation(
|
||||
default_vllm_config,
|
||||
activation_cls: type[torch.nn.Module],
|
||||
fn: object,
|
||||
op_args: tuple[str, ...],
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
x = torch.randn(num_tokens, d, dtype=dtype)
|
||||
layer = activation_cls()
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
|
||||
)
|
||||
# gelu with activation_lut_bf16 only makes sense for BF16
|
||||
if not (activation_cls is GELU and dtype != torch.bfloat16):
|
||||
raw_out = torch.empty_like(x)
|
||||
opcheck(fn, (raw_out, x, *op_args))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_gelu_tanh_and_mul(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
gate = torch.tensor(
|
||||
[
|
||||
[
|
||||
-12.0,
|
||||
-10.0,
|
||||
-9.01,
|
||||
-5.0,
|
||||
-2.0,
|
||||
-1.0,
|
||||
-0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
2.0,
|
||||
5.0,
|
||||
9.01,
|
||||
10.0,
|
||||
12.0,
|
||||
11.0,
|
||||
],
|
||||
[
|
||||
-7.5,
|
||||
-4.5,
|
||||
-3.0,
|
||||
-1.5,
|
||||
-0.75,
|
||||
-0.25,
|
||||
0.25,
|
||||
0.75,
|
||||
1.5,
|
||||
3.0,
|
||||
4.5,
|
||||
7.5,
|
||||
-11.0,
|
||||
11.0,
|
||||
8.75,
|
||||
-8.75,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
val = torch.tensor(
|
||||
[
|
||||
[
|
||||
0.25,
|
||||
-0.5,
|
||||
0.75,
|
||||
-1.0,
|
||||
1.25,
|
||||
-1.5,
|
||||
1.75,
|
||||
-2.0,
|
||||
2.25,
|
||||
-2.5,
|
||||
2.75,
|
||||
-3.0,
|
||||
3.25,
|
||||
-3.5,
|
||||
3.75,
|
||||
-4.0,
|
||||
],
|
||||
[
|
||||
-0.4,
|
||||
0.6,
|
||||
-0.8,
|
||||
1.0,
|
||||
-1.2,
|
||||
1.4,
|
||||
-1.6,
|
||||
1.8,
|
||||
-2.0,
|
||||
2.2,
|
||||
-2.4,
|
||||
2.6,
|
||||
-2.8,
|
||||
3.0,
|
||||
-3.2,
|
||||
3.4,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
x = torch.cat((val, gate), dim=-1).contiguous()
|
||||
kernel_out = torch.empty_like(val)
|
||||
torch.ops._C.gelu_tanh_and_mul(kernel_out, x)
|
||||
|
||||
torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate
|
||||
|
||||
atol = get_default_atol(kernel_out)
|
||||
rtol = get_default_rtol(kernel_out)
|
||||
torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol)
|
||||
@@ -0,0 +1,109 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the manual AllReduce + GemmaRMSNorm fusion used by MiniMax M3.
|
||||
|
||||
``fused_allreduce_gemma_rms_norm`` must match the unfused model path, i.e.
|
||||
``GemmaRMSNorm(all_reduce(partial), residual)``, both on the flashinfer fast
|
||||
path (TP>1 with flashinfer + NVSwitch) and on the eager fallback (TP==1, or when
|
||||
flashinfer is unavailable / the GPU has no NVSwitch).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.multiprocessing import spawn
|
||||
|
||||
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
|
||||
from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
|
||||
fused_allreduce_gemma_rms_norm,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@ensure_current_vllm_config()
|
||||
def _worker_fused_ar_norm(
|
||||
local_rank,
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
):
|
||||
"""Per-rank worker: compare the fused helper vs all_reduce + GemmaRMSNorm."""
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(
|
||||
world_size, 1, local_rank, port, local_rank=local_rank
|
||||
)
|
||||
|
||||
# Norm weights are identical across ranks (replicated GemmaRMSNorm).
|
||||
set_random_seed(seed)
|
||||
norm = GemmaRMSNorm(hidden_size, eps=eps).cuda().to(dtype)
|
||||
with torch.no_grad():
|
||||
norm.weight.normal_(mean=0.0, std=0.1)
|
||||
|
||||
# Residual is shared across ranks; the partial o_proj output differs per rank
|
||||
# (each rank holds a partial sum that all_reduce combines).
|
||||
torch.manual_seed(seed + 7)
|
||||
residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
torch.manual_seed(seed + 1000 + local_rank)
|
||||
partial = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Reference: the unfused model path.
|
||||
reduced = tensor_model_parallel_all_reduce(partial.clone())
|
||||
ref_out, ref_res = norm(reduced, residual.clone())
|
||||
|
||||
# Fused helper (flashinfer fast path when available, else fallback).
|
||||
out, res = fused_allreduce_gemma_rms_norm(partial.clone(), residual.clone(), norm)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=2e-2)
|
||||
torch.testing.assert_close(res, ref_res, atol=2e-2, rtol=2e-2)
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="CUDA required",
|
||||
)
|
||||
# world_size=1 exercises the TP==1 identity branch on a single GPU; >1 exercises
|
||||
# the all_reduce + GemmaRMSNorm equivalence (flashinfer kernel or fallback).
|
||||
@pytest.mark.parametrize("world_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize("num_tokens", [1, 128, 333])
|
||||
@pytest.mark.parametrize("hidden_size", [2048, 4096])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_fused_allreduce_gemma_rms_norm(
|
||||
world_size,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
eps,
|
||||
seed,
|
||||
):
|
||||
num_gpus = current_platform.device_count()
|
||||
if num_gpus < world_size:
|
||||
pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
|
||||
port = str(get_open_port())
|
||||
spawn(
|
||||
_worker_fused_ar_norm,
|
||||
args=(
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Correctness + large-token-count launch tests for fused_q_kv_rmsnorm.
|
||||
|
||||
Before the grid-dim fix the kernel used grid ``(2, num_tokens)``, which hit
|
||||
CUDA's 65535 grid-y cap for ``num_tokens >= 65536`` and failed with
|
||||
``Triton Error [CUDA]: invalid argument`` at every large chunked-prefill
|
||||
profile run. These tests pin the new grid layout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="fused_q_kv_rmsnorm requires a CUDA/ROCm device",
|
||||
)
|
||||
|
||||
|
||||
def _ref_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
x_f32 = x.to(torch.float32)
|
||||
variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
|
||||
y = x_f32 * torch.rsqrt(variance + eps) * w.to(torch.float32)
|
||||
return y.to(x.dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 17, 1024, 8192])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_fused_q_kv_rmsnorm_correctness(num_tokens: int, dtype: torch.dtype):
|
||||
torch.manual_seed(0)
|
||||
device = "cuda"
|
||||
q_size, kv_size = 192, 576
|
||||
qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
|
||||
kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
|
||||
qw = torch.randn(q_size, dtype=dtype, device=device)
|
||||
kvw = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
eps = 1e-6
|
||||
|
||||
qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
|
||||
|
||||
qr_ref = _ref_rmsnorm(qr, qw, eps)
|
||||
kv_ref = _ref_rmsnorm(kv, kvw, eps)
|
||||
|
||||
tol = dict(rtol=1e-2, atol=1e-2)
|
||||
torch.testing.assert_close(qr_out, qr_ref, **tol)
|
||||
torch.testing.assert_close(kv_out, kv_ref, **tol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [65535, 65536, 131072])
|
||||
def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
|
||||
"""Regression guard: grid used to be (2, num_tokens), hitting CUDA's
|
||||
65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
|
||||
lifts that bound to 2**31-1."""
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
q_size, kv_size = 192, 576
|
||||
qr = torch.randn(num_tokens, q_size, dtype=dtype, device=device)
|
||||
kv = torch.randn(num_tokens, kv_size, dtype=dtype, device=device)
|
||||
qw = torch.randn(q_size, dtype=dtype, device=device)
|
||||
kvw = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
|
||||
qr_out, kv_out = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-6)
|
||||
# spot-check a couple of rows against the torch reference
|
||||
for row in (0, num_tokens // 2, num_tokens - 1):
|
||||
torch.testing.assert_close(
|
||||
qr_out[row],
|
||||
_ref_rmsnorm(qr[row : row + 1], qw, 1e-6)[0],
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
kv_out[row],
|
||||
_ref_rmsnorm(kv[row : row + 1], kvw, 1e-6)[0],
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,146 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
IS_NEOX = [True, False]
|
||||
EPS_VALUES = [1e-5, 1e-6]
|
||||
SEEDS = [13]
|
||||
PARTIAL_ROPE = [True, False]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
def _apply_qk_norm_rope(
|
||||
qkv: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
q_norm: RMSNorm,
|
||||
k_norm: RMSNorm,
|
||||
rope: RotaryEmbedding,
|
||||
num_heads_q: int,
|
||||
num_heads_kv: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
q_size = num_heads_q * head_dim
|
||||
kv_size = num_heads_kv * head_dim
|
||||
|
||||
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim)
|
||||
q_by_head = q_norm.forward_native(q_by_head)
|
||||
q = q_by_head.view(q.shape)
|
||||
|
||||
k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim)
|
||||
k_by_head = k_norm.forward_native(k_by_head)
|
||||
k = k_by_head.view(k.shape)
|
||||
|
||||
q, k = rope.forward_native(positions, q, k)
|
||||
return torch.cat([q, k, v], dim=-1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="fused_qk_norm_rope custom op requires cuda and rocm platform",
|
||||
)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX)
|
||||
@pytest.mark.parametrize("eps", EPS_VALUES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
|
||||
@torch.inference_mode()
|
||||
def test_fused_qk_norm_rope_matches_reference(
|
||||
default_vllm_config,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
is_neox: bool,
|
||||
eps: float,
|
||||
seed: int,
|
||||
rotary_ratio: float,
|
||||
):
|
||||
torch.set_default_device(device)
|
||||
set_random_seed(seed)
|
||||
num_heads, num_kv_heads, head_dim = 16, 4, 128
|
||||
num_tokens = 4
|
||||
|
||||
total_dim = (num_heads + 2 * num_kv_heads) * head_dim
|
||||
qkv_base = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
|
||||
qkv_fused = qkv_base.clone()
|
||||
positions = torch.arange(num_tokens, dtype=torch.long, device=device)
|
||||
|
||||
q_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
|
||||
k_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
|
||||
q_norm.weight.data.normal_(mean=1.0, std=0.1)
|
||||
k_norm.weight.data.normal_(mean=1.0, std=0.1)
|
||||
q_weight = q_norm.weight.data
|
||||
k_weight = k_norm.weight.data
|
||||
rotary_dim = int(head_dim * rotary_ratio)
|
||||
rope = RotaryEmbedding(
|
||||
head_size=head_dim,
|
||||
rotary_dim=rotary_dim,
|
||||
max_position_embeddings=4096,
|
||||
base=10000.0,
|
||||
is_neox_style=is_neox,
|
||||
dtype=dtype,
|
||||
).to(device)
|
||||
|
||||
ref_result = _apply_qk_norm_rope(
|
||||
qkv=qkv_base,
|
||||
positions=positions,
|
||||
q_norm=q_norm,
|
||||
k_norm=k_norm,
|
||||
rope=rope,
|
||||
num_heads_q=num_heads,
|
||||
num_heads_kv=num_kv_heads,
|
||||
head_dim=head_dim,
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.fused_qk_norm_rope,
|
||||
(
|
||||
qkv_fused.clone(),
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
),
|
||||
)
|
||||
|
||||
torch.ops._C.fused_qk_norm_rope(
|
||||
qkv_fused,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope.cos_sin_cache,
|
||||
is_neox,
|
||||
positions.view(-1),
|
||||
)
|
||||
|
||||
if dtype == torch.float16:
|
||||
ATOL, RTOL = (2e-3, 2e-3)
|
||||
else:
|
||||
ATOL, RTOL = (1e-2, 1e-2)
|
||||
|
||||
torch.testing.assert_close(
|
||||
qkv_fused,
|
||||
ref_result,
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
@@ -0,0 +1,335 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.utils import fp8_ulp_distance, opcheck
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_group_quant_int8,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
|
||||
*[(4096, i) for i in [1, 64, 5137]],
|
||||
]
|
||||
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [None, [1, 64], [1, 128]]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
EPS = 1e-6
|
||||
|
||||
## Helpers
|
||||
|
||||
|
||||
def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor:
|
||||
return torch.as_tensor(x, dtype=torch.float32, device="cuda")
|
||||
|
||||
|
||||
def ref_rms_norm(
|
||||
rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor | None
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
out, residual = rms_norm_layer.forward_native(x, residual)
|
||||
else:
|
||||
out = rms_norm_layer.forward_native(x)
|
||||
|
||||
return out, residual
|
||||
|
||||
|
||||
def ref_dynamic_per_token_or_block_quant(
|
||||
rms_norm_layer: RMSNorm,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if scale_ub is not None:
|
||||
assert quant_dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Norm
|
||||
torch_out, residual = ref_rms_norm(rms_norm_layer, x, residual)
|
||||
|
||||
# Quant
|
||||
if group_size is not None:
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
torch_out, scales = per_token_group_quant_fp8(
|
||||
torch_out, group_size=group_size[1], use_ue8m0=False
|
||||
)
|
||||
else:
|
||||
assert quant_dtype == torch.int8
|
||||
torch_out, scales = per_token_group_quant_int8(
|
||||
torch_out, group_size=group_size[1]
|
||||
)
|
||||
else:
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
torch_out, scales = ops.scaled_fp8_quant(
|
||||
torch_out, scale_ub=scale_ub, use_per_token_if_dynamic=True
|
||||
)
|
||||
else:
|
||||
assert quant_dtype == torch.int8
|
||||
torch_out, scales, _ = ops.scaled_int8_quant(torch_out)
|
||||
|
||||
return torch_out, scales, residual
|
||||
|
||||
|
||||
def ref_impl(
|
||||
rms_norm_layer: RMSNorm,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ref_dynamic_per_token_or_block_quant(
|
||||
rms_norm_layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
|
||||
|
||||
def ops_dynamic_per_token_or_block_quant(
|
||||
weight: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
if group_size is not None:
|
||||
out, scales = ops.rms_norm_per_block_quant(
|
||||
x,
|
||||
weight,
|
||||
EPS,
|
||||
quant_dtype,
|
||||
group_size,
|
||||
scale_ub,
|
||||
residual,
|
||||
True,
|
||||
tma_alignment,
|
||||
)
|
||||
scales = scales.contiguous()
|
||||
else:
|
||||
out, scales = ops.rms_norm_dynamic_per_token_quant(
|
||||
x, weight, EPS, quant_dtype, scale_ub, residual
|
||||
)
|
||||
return out, scales, residual
|
||||
|
||||
|
||||
def ops_impl(
|
||||
weight: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ops_dynamic_per_token_or_block_quant(
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
if group_size is not None and hidden_size % group_size[1] != 0:
|
||||
# skip
|
||||
pytest.skip("Skip non-divisible group sizes")
|
||||
|
||||
if group_size is not None and has_scale_ub:
|
||||
# blockwise baseline doesn't support scale_ub
|
||||
pytest.skip("scale_ub not supported for blockwise/group quantization")
|
||||
|
||||
if (
|
||||
group_size is None or quant_dtype != current_platform.fp8_dtype()
|
||||
) and tma_alignment != 0:
|
||||
# TMA alignment is only supported for groupwise fp8 kernels
|
||||
pytest.skip("tma alignment not supported for per-token or int8 quantization")
|
||||
|
||||
if (
|
||||
group_size is not None
|
||||
and tma_alignment != 0
|
||||
and hidden_size // group_size[1] % tma_alignment == 0
|
||||
):
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
pytest.skip("Skip TMA alignment cases where no extra padding is added")
|
||||
|
||||
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
|
||||
# skip
|
||||
pytest.skip("scale_ub only supported for fp8 quantization")
|
||||
|
||||
layer = RMSNorm(hidden_size, EPS).to(dtype=dtype)
|
||||
|
||||
# Make weights
|
||||
layer.weight.data.normal_(mean=1.0, std=0.1)
|
||||
|
||||
# Make inputs: use a wider tensor and slice to create a non-contiguous
|
||||
# (strided) input when strided_input=True. The last dimension stride
|
||||
# remains 1, which the kernel requires.
|
||||
scale = 1 / (hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x = torch.randn(num_tokens, last_dim, dtype=dtype) * scale
|
||||
x = x[:, :hidden_size]
|
||||
|
||||
# dim 1 gets special-cased
|
||||
x_is_strided = strided_input and num_tokens != 1
|
||||
# check that the input is strided iff we expect it to be
|
||||
assert x.is_contiguous() != x_is_strided
|
||||
|
||||
# Residual must still be contiguous
|
||||
residual = (
|
||||
torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
|
||||
if add_residual
|
||||
else None
|
||||
)
|
||||
if has_scale_ub:
|
||||
rms_x, _ = ref_rms_norm(layer, x, residual)
|
||||
scale_ub = torch.mean(rms_x).to(dtype=torch.float32, device="cuda")
|
||||
else:
|
||||
scale_ub = None
|
||||
|
||||
ref_out, ref_scales, ref_residual = ref_impl(
|
||||
layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
ops_out, ops_scales, ops_residual = ops_impl(
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
assert ref_out.dtype == quant_dtype
|
||||
assert ops_out.dtype == quant_dtype
|
||||
|
||||
# Per-block bf16 scales: allow a small relative tolerance for a few groups
|
||||
# whose abs-max flips by one ULP between the fused and reference paths. The
|
||||
# per-token and fp32 paths stay strict.
|
||||
relax_block_rocm = (
|
||||
group_size is not None
|
||||
and dtype == torch.bfloat16
|
||||
and current_platform.is_rocm()
|
||||
)
|
||||
|
||||
def scales_close(rtol: float, atol: float) -> bool:
|
||||
if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
|
||||
return True
|
||||
return relax_block_rocm and torch.allclose(
|
||||
ref_scales, ops_scales, rtol=1e-2, atol=atol
|
||||
)
|
||||
|
||||
if quant_dtype == torch.int8:
|
||||
assert scales_close(rtol=1e-5, atol=1e-6)
|
||||
# big atol to account for round-off errors.
|
||||
assert torch.allclose(ref_out, ops_out, atol=1)
|
||||
else:
|
||||
assert scales_close(rtol=1e-5, atol=1e-8)
|
||||
a = ref_out.to(dtype=torch.float32)
|
||||
b = ops_out.to(dtype=torch.float32)
|
||||
ok = torch.allclose(a, b, atol=1e-6)
|
||||
if not ok:
|
||||
if relax_block_rocm:
|
||||
# ULP-flipped group scale can cross an E4M3 tie; tolerate a
|
||||
# bounded count of isolated fp8 outliers.
|
||||
ulp = fp8_ulp_distance(ref_out, ops_out)
|
||||
max_outliers = ulp.numel() // 100_000 + 8
|
||||
ok = int((ulp > 0).sum().item()) <= max_outliers
|
||||
else:
|
||||
# CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
|
||||
if group_size is None:
|
||||
a_deq = a * ref_scales.view(-1, 1)
|
||||
b_deq = b * ops_scales.view(-1, 1)
|
||||
else:
|
||||
a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
|
||||
b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
|
||||
# NOTE: It is possible that some future test cases trigger this
|
||||
# max diff due to precision issues. If such an error is
|
||||
# encountered, it's recommended to inspect the differences between
|
||||
# all corresponding elements from each tensor (e.g. by looping over
|
||||
# them) and checking how many the max diff error shows up on (just
|
||||
# a few bad elements should still be considered acceptable).
|
||||
ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
|
||||
assert ok
|
||||
if add_residual:
|
||||
assert torch.allclose(ref_residual, ops_residual)
|
||||
|
||||
output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
|
||||
if group_size is None:
|
||||
scales = torch.empty(
|
||||
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
|
||||
)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_dynamic_per_token_quant,
|
||||
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
|
||||
)
|
||||
else:
|
||||
assert hidden_size % group_size[1] == 0
|
||||
num_groups = hidden_size // group_size[1]
|
||||
scales = torch.empty(
|
||||
(num_groups, num_tokens),
|
||||
device=x.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_per_block_quant,
|
||||
(
|
||||
output,
|
||||
x,
|
||||
layer.weight,
|
||||
scales,
|
||||
1e-5,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size[1],
|
||||
True, # is_scale_transposed
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,103 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Tests that FusedRMSNormGated decomposes correctly under torch.compile,
|
||||
matching the eager triton kernel output."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16]
|
||||
HIDDEN_SIZES = [128, 512]
|
||||
NUM_TOKENS = [64, 128]
|
||||
ACTIVATIONS = ["swish", "sigmoid"]
|
||||
ELEMENTWISE_AFFINE = [True, False]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("activation", ACTIVATIONS)
|
||||
@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_compiled_vs_eager(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
activation: str,
|
||||
elementwise_affine: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""forward_native decomposition matches forward_cuda triton kernel."""
|
||||
torch._dynamo.reset()
|
||||
set_random_seed(seed)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
module = FusedRMSNormGated(
|
||||
hidden_size,
|
||||
elementwise_affine=elementwise_affine,
|
||||
eps=1e-5,
|
||||
activation=activation,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
g = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# forward_cuda may modify x in-place, so clone inputs
|
||||
cuda_out = module.forward_cuda(x.clone(), g.clone())
|
||||
compiled_native = torch.compile(module.forward_native, fullgraph=True)
|
||||
native_out = compiled_native(x.clone(), g.clone())
|
||||
|
||||
torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 16, 32, 128),
|
||||
(2, 8, 16, 64),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("activation", ACTIVATIONS)
|
||||
@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_compiled_vs_eager_multidim(
|
||||
default_vllm_config,
|
||||
shape: tuple,
|
||||
activation: str,
|
||||
elementwise_affine: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""forward_native decomposition handles multi-dimensional inputs."""
|
||||
torch._dynamo.reset()
|
||||
set_random_seed(seed)
|
||||
device = torch.device("cuda:0")
|
||||
head_dim = shape[-1]
|
||||
|
||||
module = FusedRMSNormGated(
|
||||
head_dim,
|
||||
elementwise_affine=elementwise_affine,
|
||||
eps=1e-5,
|
||||
activation=activation,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
x = torch.randn(*shape, dtype=dtype, device=device)
|
||||
g = torch.randn(*shape, dtype=dtype, device=device)
|
||||
|
||||
# forward_cuda may modify x in-place, so clone inputs
|
||||
cuda_out = module.forward_cuda(x.clone(), g.clone())
|
||||
compiled_native = torch.compile(module.forward_native, fullgraph=True)
|
||||
native_out = compiled_native(x.clone(), g.clone())
|
||||
|
||||
torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_group_quant_int8,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
QUANT_DTYPES = [current_platform.fp8_dtype(), torch.int8]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [64, *VEC_HIDDEN_SIZES, 2048, 5120]],
|
||||
*[(16, i) for i in [64, *VEC_HIDDEN_SIZES, 5120]],
|
||||
*[(128, i) for i in [64, *VEC_HIDDEN_SIZES]],
|
||||
*[(512, i) for i in [64, 5120]],
|
||||
]
|
||||
SCALE_UBS = [False]
|
||||
GROUP_SIZES = [64, 128]
|
||||
IS_SCALE_TRANSPOSED = [False, True]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [i for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
|
||||
|
||||
|
||||
def ref_silu_and_mul_per_block_quant(
|
||||
x: torch.Tensor,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reference implementation: unfused SiLU+Mul then group quantization."""
|
||||
hidden = x.shape[-1] // 2
|
||||
gate, up = x.split(hidden, dim=-1)
|
||||
silu_out = F.silu(gate) * up
|
||||
|
||||
if quant_dtype == current_platform.fp8_dtype():
|
||||
return per_token_group_quant_fp8(
|
||||
silu_out, group_size=group_size, use_ue8m0=False
|
||||
)
|
||||
elif quant_dtype == torch.int8:
|
||||
return per_token_group_quant_int8(silu_out, group_size=group_size)
|
||||
else:
|
||||
raise ValueError(f"Unsupported quant_dtype: {quant_dtype}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("group_size", GROUP_SIZES)
|
||||
@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device_idx", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_and_mul_per_block_quant(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: int,
|
||||
is_scale_transposed: bool,
|
||||
seed: int,
|
||||
device_idx: str,
|
||||
) -> None:
|
||||
"""Test SiLU+Mul+Block Quantization kernel correctness."""
|
||||
torch.accelerator.set_device_index(device_idx)
|
||||
device = f"cuda:{device_idx}"
|
||||
torch.random.manual_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
if hidden_size % group_size != 0:
|
||||
return
|
||||
|
||||
if has_scale_ub:
|
||||
pytest.skip("Scale upper bound not yet supported")
|
||||
|
||||
scale = 1 / hidden_size
|
||||
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
|
||||
|
||||
# Reference implementation
|
||||
ref_out, ref_scales = ref_silu_and_mul_per_block_quant(x, quant_dtype, group_size)
|
||||
|
||||
# Fused kernel implementation
|
||||
ops_out, ops_scales = ops.silu_and_mul_per_block_quant(
|
||||
x, group_size, quant_dtype, None, is_scale_transposed
|
||||
)
|
||||
|
||||
# Check for NaN/Inf
|
||||
assert not torch.isnan(ops_out.float()).any(), "Kernel output contains NaN"
|
||||
assert not torch.isinf(ops_out.float()).any(), "Kernel output contains Inf"
|
||||
assert not torch.isnan(ops_scales).any(), "Kernel scales contain NaN"
|
||||
assert not torch.isinf(ops_scales).any(), "Kernel scales contain Inf"
|
||||
|
||||
# Check dtypes
|
||||
assert ref_out.dtype == quant_dtype
|
||||
assert ops_out.dtype == quant_dtype
|
||||
|
||||
# Check scales match
|
||||
torch.testing.assert_close(ref_scales, ops_scales, rtol=1e-5, atol=1e-5)
|
||||
|
||||
# Check output correctness via dequantized values
|
||||
ref_scales_expanded = ref_scales.repeat_interleave(group_size, dim=1)
|
||||
ops_scales_expanded = ops_scales.repeat_interleave(group_size, dim=1)
|
||||
ref_deq = ref_out.to(dtype=torch.float32) * ref_scales_expanded
|
||||
ops_deq = ops_out.to(dtype=torch.float32) * ops_scales_expanded
|
||||
torch.testing.assert_close(ref_deq, ops_deq, atol=5e-2, rtol=5e-2)
|
||||
|
||||
# opcheck
|
||||
output = torch.empty(num_tokens, hidden_size, device=device, dtype=quant_dtype)
|
||||
num_groups = hidden_size // group_size
|
||||
if is_scale_transposed:
|
||||
scales = torch.empty(num_groups, num_tokens, device=device, dtype=torch.float32)
|
||||
else:
|
||||
scales = torch.empty(num_tokens, num_groups, device=device, dtype=torch.float32)
|
||||
opcheck(
|
||||
torch.ops._C.silu_and_mul_per_block_quant,
|
||||
(output, x, scales, group_size, None, is_scale_transposed),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16])
|
||||
@pytest.mark.parametrize("hidden_size", [4096])
|
||||
@pytest.mark.parametrize("num_tokens", [128])
|
||||
@pytest.mark.parametrize("group_size", [128])
|
||||
def test_silu_block_quant_shapes(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
hidden_size: int,
|
||||
num_tokens: int,
|
||||
group_size: int,
|
||||
):
|
||||
"""Test that output shapes are correct."""
|
||||
torch.set_default_device("cuda")
|
||||
x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device="cuda")
|
||||
|
||||
# Row-major scales
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=group_size,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=False,
|
||||
)
|
||||
assert out.shape == (num_tokens, hidden_size)
|
||||
assert scales.shape == (num_tokens, hidden_size // group_size)
|
||||
|
||||
# Column-major scales (logical shape same after .t() in _custom_ops)
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=group_size,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=True,
|
||||
)
|
||||
assert out.shape == (num_tokens, hidden_size)
|
||||
assert scales.shape == (num_tokens, hidden_size // group_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16])
|
||||
@pytest.mark.parametrize("batch_size", [1, 16, 256])
|
||||
@pytest.mark.parametrize("hidden_size", [1024, 5120, 14336])
|
||||
def test_silu_block_quant_edge_cases(
|
||||
default_vllm_config, dtype: torch.dtype, batch_size: int, hidden_size: int
|
||||
):
|
||||
"""Test edge cases: single token, large batch, large hidden size."""
|
||||
torch.set_default_device("cuda")
|
||||
x = torch.randn(batch_size, hidden_size * 2, dtype=dtype, device="cuda")
|
||||
|
||||
out, scales = ops.silu_and_mul_per_block_quant(
|
||||
x,
|
||||
group_size=128,
|
||||
quant_dtype=current_platform.fp8_dtype(),
|
||||
is_scale_transposed=False,
|
||||
)
|
||||
|
||||
assert out.shape == (batch_size, hidden_size)
|
||||
assert out.dtype == current_platform.fp8_dtype()
|
||||
assert scales.dtype == torch.float32
|
||||
assert not torch.isnan(out.float()).any()
|
||||
assert not torch.isnan(scales).any()
|
||||
assert not torch.isinf(scales).any()
|
||||
@@ -0,0 +1,250 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from tests.kernels.utils import fp8_ulp_distance, opcheck
|
||||
from vllm import ir
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx90a
|
||||
|
||||
on_mi250 = on_gfx90a()
|
||||
else:
|
||||
on_mi250 = False
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16, torch.float]
|
||||
NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing
|
||||
HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing
|
||||
ADD_RESIDUAL = [False, True] if not on_mi250 else [True]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
def _rms_norm_tolerance(dtype: torch.dtype) -> dict[str, float]:
|
||||
return ir.ops.rms_norm.get_tolerance(dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
layer = RMSNorm(hidden_size).to(dtype=dtype)
|
||||
layer.weight.data.normal_(mean=1.0, std=0.1)
|
||||
scale = 1 / (2 * hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x = torch.randn(num_tokens, last_dim, dtype=dtype)
|
||||
x = x[..., :hidden_size]
|
||||
assert x.is_contiguous() != strided_input
|
||||
x *= scale
|
||||
residual = torch.randn_like(x) * scale if add_residual else None
|
||||
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_out = layer.forward_native(x, residual)
|
||||
out = layer(x, residual)
|
||||
# NOTE(woosuk): LayerNorm operators (including RMS) typically have larger
|
||||
# numerical errors than other operators because they involve reductions.
|
||||
# Therefore, we use a larger tolerance.
|
||||
if add_residual:
|
||||
torch.testing.assert_close(out[0], ref_out[0], atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(out[1], ref_out[1], atol=1e-2, rtol=1e-2)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
|
||||
if residual is not None:
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm,
|
||||
(x, residual, layer.weight.data, layer.variance_epsilon),
|
||||
)
|
||||
else:
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm, (out, x, layer.weight.data, layer.variance_epsilon)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_rms_norm_weightless(
|
||||
default_vllm_config,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
layer = RMSNorm(hidden_size, has_weight=False).to(dtype=dtype)
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
residual = torch.randn_like(x) if add_residual else None
|
||||
|
||||
ref_out = layer.forward_native(x, residual)
|
||||
out = layer(x, residual)
|
||||
tol = _rms_norm_tolerance(dtype)
|
||||
if add_residual:
|
||||
torch.testing.assert_close(out[0], ref_out[0], **tol)
|
||||
torch.testing.assert_close(out[1], ref_out[1], **tol)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, **tol)
|
||||
|
||||
if residual is not None:
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm,
|
||||
(x, residual, None, layer.variance_epsilon),
|
||||
)
|
||||
else:
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm,
|
||||
(out, x, None, layer.variance_epsilon),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_scale", [0.01, 1.0, 10.0])
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("strided_input", [False, True])
|
||||
def test_fused_rms_norm_quant(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_scale: float,
|
||||
seed: int,
|
||||
device: str,
|
||||
strided_input: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
weight = torch.empty(hidden_size, dtype=dtype).normal_(mean=1.0, std=0.1)
|
||||
scale = 1 / (2 * hidden_size)
|
||||
last_dim = 2 * hidden_size if strided_input else hidden_size
|
||||
x_base = torch.randn(num_tokens, last_dim, dtype=dtype)
|
||||
x = x_base[..., :hidden_size]
|
||||
assert x.is_contiguous() != strided_input
|
||||
|
||||
x *= scale
|
||||
if add_residual:
|
||||
residual = torch.randn_like(x) * scale
|
||||
residual_fused = residual.clone()
|
||||
else:
|
||||
residual = residual_fused = None
|
||||
|
||||
out_norm = torch.empty_like(x)
|
||||
out_quant = torch.empty_like(x, dtype=FP8_DTYPE)
|
||||
out_quant_fused = torch.empty_like(out_quant)
|
||||
|
||||
quant_scale_t = torch.tensor(quant_scale, dtype=torch.float32)
|
||||
|
||||
if add_residual:
|
||||
torch.ops._C.fused_add_rms_norm_static_fp8_quant(
|
||||
out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6
|
||||
)
|
||||
|
||||
# Unfused kernel is in-place so it goes second
|
||||
# Also use a separate clone of x to avoid modifying the input
|
||||
x_unfused_base = x_base.clone()
|
||||
x_unfused = x_unfused_base[..., :hidden_size]
|
||||
assert x_unfused.is_contiguous() != strided_input
|
||||
torch.ops._C.fused_add_rms_norm(x_unfused, residual, weight, 1e-6)
|
||||
torch.ops._C.static_scaled_fp8_quant(
|
||||
out_quant, x_unfused.contiguous(), quant_scale_t
|
||||
)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(residual_fused, residual, atol=1e-2, rtol=1e-2)
|
||||
opcheck(
|
||||
torch.ops._C.fused_add_rms_norm_static_fp8_quant,
|
||||
(out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6),
|
||||
)
|
||||
else:
|
||||
torch.ops._C.rms_norm_static_fp8_quant(
|
||||
out_quant_fused, x, weight, quant_scale_t, 1e-6
|
||||
)
|
||||
|
||||
torch.ops._C.rms_norm(out_norm, x, weight, 1e-6)
|
||||
torch.ops._C.static_scaled_fp8_quant(out_quant, out_norm, quant_scale_t)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_static_fp8_quant,
|
||||
(out_quant_fused, x, weight, quant_scale_t, 1e-6),
|
||||
)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# Fused and unfused FP8 paths can land on opposite sides of an E4M3 tie;
|
||||
# tolerate a tiny number of isolated fp8 outliers on ROCm.
|
||||
ulp = fp8_ulp_distance(out_quant, out_quant_fused)
|
||||
max_outliers = ulp.numel() // 100_000 + 8
|
||||
num_outliers = int((ulp > 0).sum().item())
|
||||
assert num_outliers <= max_outliers, (
|
||||
f"FP8 quant mismatch: {num_outliers} fp8 outliers (allowed {max_outliers})"
|
||||
)
|
||||
else:
|
||||
torch.testing.assert_close(
|
||||
out_quant.to(dtype=torch.float32),
|
||||
out_quant_fused.to(dtype=torch.float32),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
|
||||
device = CUDA_DEVICES[0]
|
||||
torch.set_default_device(device)
|
||||
|
||||
num_tokens, hidden_size = 32, 1024
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device)
|
||||
layer.weight.data.normal_(mean=0.0, std=0.1)
|
||||
|
||||
# Gemma uses fp32 weight parameter while activations can be bf16.
|
||||
assert layer.weight.dtype == torch.float32
|
||||
out = layer(x)
|
||||
|
||||
x_fp32 = x.float()
|
||||
weight_fp32 = layer.weight.data.float() + 1.0
|
||||
variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
|
||||
ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to(
|
||||
x.dtype
|
||||
)
|
||||
|
||||
assert out.dtype == x.dtype
|
||||
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.multiprocessing import spawn
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.model_executor.layers.minimax_rms_norm import (
|
||||
MiniMaxText01RMSNormTP,
|
||||
rms_norm_tp,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@ensure_current_vllm_config()
|
||||
def _worker_forward_qk(
|
||||
local_rank,
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_q_full,
|
||||
hidden_k_full,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
):
|
||||
"""Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel."""
|
||||
|
||||
if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
|
||||
cleanup_dist_env_and_memory()
|
||||
return
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(
|
||||
world_size, 1, local_rank, port, local_rank=local_rank
|
||||
)
|
||||
|
||||
hq = hidden_q_full // world_size
|
||||
hk = hidden_k_full // world_size
|
||||
|
||||
q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda()
|
||||
k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda()
|
||||
|
||||
set_random_seed(seed)
|
||||
qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda")
|
||||
kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda")
|
||||
q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq])
|
||||
k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk])
|
||||
|
||||
torch.manual_seed(seed + 1000 + local_rank)
|
||||
qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
|
||||
|
||||
# Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
|
||||
# the variance (it is the tp==1 / already-reduced building block), so the
|
||||
# multi-rank reference must use the eager path that performs the global
|
||||
# variance all-reduce, matching the fused kernel below.
|
||||
ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
hq,
|
||||
hk,
|
||||
world_size,
|
||||
eps,
|
||||
)
|
||||
|
||||
# Set up Lamport workspace.
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
from vllm.model_executor.layers.minimax_rms_norm.lamport_workspace import (
|
||||
get_allreduce_workspace,
|
||||
)
|
||||
|
||||
workspace = get_allreduce_workspace(
|
||||
rank=local_rank,
|
||||
world_size=world_size,
|
||||
max_tokens=num_tokens,
|
||||
process_group=get_tp_group().cpu_group,
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.minimax_allreduce_rms_qk,
|
||||
(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
workspace,
|
||||
hq,
|
||||
hk,
|
||||
local_rank,
|
||||
world_size,
|
||||
eps,
|
||||
),
|
||||
)
|
||||
fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk(
|
||||
qkv.clone(),
|
||||
q_norm.weight,
|
||||
k_norm.weight,
|
||||
workspace,
|
||||
hq,
|
||||
hk,
|
||||
local_rank,
|
||||
world_size,
|
||||
eps,
|
||||
)
|
||||
_, _, fused_v = qkv.split([hq, hk, hk], dim=-1)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
fused_q,
|
||||
ref_q,
|
||||
atol=3e-2,
|
||||
rtol=3e-2,
|
||||
)
|
||||
torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2)
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="CUDA required",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2, 4, 8])
|
||||
@pytest.mark.parametrize("num_tokens", [1, 128, 333])
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_dims",
|
||||
[(6144, 1024)],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_minimax_reduce_rms_qk(
|
||||
world_size,
|
||||
num_tokens,
|
||||
hidden_dims,
|
||||
dtype,
|
||||
eps,
|
||||
seed,
|
||||
):
|
||||
num_gpus = current_platform.device_count()
|
||||
if num_gpus < world_size:
|
||||
pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
|
||||
hidden_q_full, hidden_k_full = hidden_dims
|
||||
port = str(get_open_port())
|
||||
spawn(
|
||||
_worker_forward_qk,
|
||||
args=(
|
||||
world_size,
|
||||
port,
|
||||
num_tokens,
|
||||
hidden_q_full,
|
||||
hidden_k_full,
|
||||
dtype,
|
||||
seed,
|
||||
eps,
|
||||
),
|
||||
nprocs=world_size,
|
||||
join=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda() or not HAS_TRITON,
|
||||
reason="CUDA and Triton required",
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
|
||||
@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("tp_world", [1, 4, 8])
|
||||
@pytest.mark.parametrize("eps", [1e-6])
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
def test_minimax_qk_norm_triton_fallback(
|
||||
monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
|
||||
):
|
||||
"""Single-GPU check: Triton fallback kernels vs the pure-torch reference.
|
||||
|
||||
The all-reduce is a TP communication barrier, so it is monkeypatched to
|
||||
identity here; both the Triton path and the reference see the same
|
||||
(patched) reduction. This validates the kernel math and the folded
|
||||
``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
|
||||
are the per-rank q/k segment widths.
|
||||
"""
|
||||
monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
|
||||
|
||||
q_size, kv_size = hidden_dims
|
||||
device = "cuda"
|
||||
torch.manual_seed(seed)
|
||||
qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
|
||||
q_weight = torch.randn(q_size, dtype=dtype, device=device)
|
||||
k_weight = torch.randn(kv_size, dtype=dtype, device=device)
|
||||
|
||||
q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
|
||||
qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
|
||||
)
|
||||
q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
|
||||
qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
|
||||
)
|
||||
|
||||
torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
|
||||
torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def generate_test_data(
|
||||
num_tokens: int,
|
||||
num_q_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
max_position_embeddings: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
):
|
||||
"""Generate test data for given configuration."""
|
||||
set_random_seed(42)
|
||||
# Create 2D positions (3, num_tokens) for multimodal case
|
||||
positions = torch.randint(
|
||||
0, max_position_embeddings // 4, (3, num_tokens), device=device
|
||||
)
|
||||
|
||||
# Create query and key tensors
|
||||
query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
|
||||
key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
|
||||
|
||||
return positions, query, key
|
||||
|
||||
|
||||
class MRoPETestInfo(NamedTuple):
|
||||
model_name: str
|
||||
# https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317
|
||||
atol: float = 1e-2
|
||||
rtol: float = 1.6e-2
|
||||
marks: list[pytest.MarkDecorator] = []
|
||||
|
||||
|
||||
MODELS_TO_TEST = [
|
||||
MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen3-VL-4B-Instruct"),
|
||||
MRoPETestInfo(model_name="Qwen/Qwen3-VL-30B-A3B-Instruct"),
|
||||
]
|
||||
|
||||
num_tokens_list = [11, 8192]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"model_info, model_name",
|
||||
[
|
||||
pytest.param(test_config, test_config.model_name, marks=test_config.marks)
|
||||
for test_config in MODELS_TO_TEST
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [1, 2])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_tokens", num_tokens_list)
|
||||
def test_mrope(
|
||||
default_vllm_config,
|
||||
model_name: str,
|
||||
model_info: MRoPETestInfo,
|
||||
tp_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_tokens: int,
|
||||
):
|
||||
atol = model_info.atol
|
||||
rtol = model_info.rtol
|
||||
|
||||
config = get_config(model_name, False).get_text_config()
|
||||
|
||||
# get the model config
|
||||
total_num_kv_heads = config.num_key_value_heads
|
||||
total_num_heads = config.num_attention_heads
|
||||
num_heads = total_num_heads // tp_size
|
||||
num_kv_heads = max(1, total_num_kv_heads // tp_size)
|
||||
head_dim = (
|
||||
config.head_dim
|
||||
if hasattr(config, "head_dim")
|
||||
else config.hidden_size // total_num_heads
|
||||
)
|
||||
is_neox_style = True
|
||||
|
||||
max_position = config.max_position_embeddings
|
||||
|
||||
mrope_helper_class = get_rope(
|
||||
head_size=head_dim,
|
||||
max_position=max_position,
|
||||
is_neox_style=is_neox_style,
|
||||
rope_parameters=config.rope_parameters,
|
||||
dtype=dtype,
|
||||
).to(device=device)
|
||||
|
||||
# create q k v input tensors
|
||||
# create rotary pos emb input tensors
|
||||
positions, query, key = generate_test_data(
|
||||
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
|
||||
)
|
||||
|
||||
query_native, key_native = mrope_helper_class.forward_native(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
query_cuda, key_cuda = mrope_helper_class.forward_cuda(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"model_info, model_name",
|
||||
[
|
||||
pytest.param(test_config, test_config.model_name, marks=test_config.marks)
|
||||
for test_config in MODELS_TO_TEST
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [1, 2])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_tokens", num_tokens_list)
|
||||
def test_mrope_torch_compile_tracing(
|
||||
default_vllm_config,
|
||||
model_name: str,
|
||||
model_info: MRoPETestInfo,
|
||||
tp_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_tokens: int,
|
||||
):
|
||||
atol = model_info.atol
|
||||
rtol = model_info.rtol
|
||||
|
||||
config = get_config(model_name, False).get_text_config()
|
||||
|
||||
# get the model config
|
||||
total_num_kv_heads = config.num_key_value_heads
|
||||
total_num_heads = config.num_attention_heads
|
||||
num_heads = total_num_heads // tp_size
|
||||
num_kv_heads = max(1, total_num_kv_heads // tp_size)
|
||||
head_dim = (
|
||||
config.head_dim
|
||||
if hasattr(config, "head_dim")
|
||||
else config.hidden_size // total_num_heads
|
||||
)
|
||||
is_neox_style = True
|
||||
max_position = config.max_position_embeddings
|
||||
|
||||
mrope_helper_class = get_rope(
|
||||
head_size=head_dim,
|
||||
max_position=max_position,
|
||||
is_neox_style=is_neox_style,
|
||||
rope_parameters=config.rope_parameters,
|
||||
dtype=dtype,
|
||||
).to(device=device)
|
||||
|
||||
# Generate test data
|
||||
positions, query, key = generate_test_data(
|
||||
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
|
||||
)
|
||||
|
||||
# Create a wrapper that makes the in-place function appear functional
|
||||
def functional_forward_cuda(pos, q, k):
|
||||
"""Wrapper that converts in-place operation to functional style
|
||||
|
||||
CUDA Graph does not support in-place operations.
|
||||
This wrapper creates working copies of the
|
||||
input tensors and modifies them.
|
||||
"""
|
||||
q_work = q.clone() # Create working copies
|
||||
k_work = k.clone()
|
||||
# Your in-place function modifies q_work and k_work
|
||||
mrope_helper_class.forward_cuda(pos, q_work, k_work)
|
||||
return q_work, k_work # Return the modified tensors
|
||||
|
||||
# Get reference results
|
||||
query_native, key_native = mrope_helper_class.forward_native(
|
||||
positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
|
||||
try:
|
||||
compiled_forward_cuda = torch.compile(
|
||||
functional_forward_cuda,
|
||||
fullgraph=True,
|
||||
backend="inductor",
|
||||
mode="reduce-overhead",
|
||||
dynamic=False,
|
||||
)
|
||||
|
||||
# Run compiled version
|
||||
query_compiled_cuda, key_compiled_cuda = compiled_forward_cuda(
|
||||
positions,
|
||||
query,
|
||||
key,
|
||||
)
|
||||
|
||||
# Run original version for comparison
|
||||
query_cuda = query.clone()
|
||||
key_cuda = key.clone()
|
||||
mrope_helper_class.forward_cuda(positions, query_cuda, key_cuda)
|
||||
|
||||
# Verify results
|
||||
torch.testing.assert_close(
|
||||
query_compiled_cuda, query_cuda, atol=atol, rtol=rtol
|
||||
)
|
||||
torch.testing.assert_close(key_compiled_cuda, key_cuda, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(
|
||||
query_compiled_cuda, query_native, atol=atol, rtol=rtol
|
||||
)
|
||||
torch.testing.assert_close(key_compiled_cuda, key_native, atol=atol, rtol=rtol)
|
||||
|
||||
print("✓ forward_cuda successfully traced with torch.compile inductor")
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"forward_cuda failed to trace with torch.compile inductor: {e}")
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for miscellaneous utilities
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
|
||||
|
||||
def test_convert_fp8_opcheck():
|
||||
data = torch.randn((256, 256), dtype=torch.float32, device="cuda")
|
||||
result = torch.empty_like(data, dtype=torch.float8_e4m3fn)
|
||||
opcheck(torch.ops._C_cache_ops.convert_fp8, (result, data, 1.0, "fp8"))
|
||||
|
||||
|
||||
# TODO: Add this back, currently fails with
|
||||
# csrc/cuda_utils_kernels.cu:15 'invalid argument'
|
||||
# @pytest.mark.skipif(not current_platform.is_cuda(),
|
||||
# reason="Only supported for CUDA")
|
||||
# def test_cuda_utils_opcheck():
|
||||
# opcheck(torch.ops._C_cuda_utils.get_device_attribute, (0, 0))
|
||||
# opcheck(
|
||||
# torch.ops._C_cuda_utils.
|
||||
# get_max_shared_memory_per_block_device_attribute, (0, ))
|
||||
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm._custom_ops import permute_cols
|
||||
|
||||
if not hasattr(torch.ops._C, "permute_cols"):
|
||||
pytest.skip(reason="permute_cols is not supported on ROCm", allow_module_level=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 512), (544, 4096), (67, 8192)])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
def test_permute_cols(shape, dtype):
|
||||
x = torch.randn(shape, dtype=dtype).cuda()
|
||||
perm = torch.randperm(x.shape[1]).to(torch.int).cuda()
|
||||
opcheck(torch.ops._C.permute_cols, (x, perm))
|
||||
y = permute_cols(x, perm)
|
||||
torch.testing.assert_close(y, x[:, perm])
|
||||
@@ -0,0 +1,193 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from itertools import product
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
IS_NEOX_STYLE = [True, False]
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
HEAD_SIZES = [64, 80, 120, 256]
|
||||
ROTARY_DIMS = [None, 32] # None means rotary dim == head size
|
||||
NUM_HEADS = [17] # Arbitrary values for testing
|
||||
BATCH_SIZES = [5] # Arbitrary values for testing
|
||||
SEQ_LENS = [11, 8192] # Arbitrary values for testing
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
USE_KEY = [True, False]
|
||||
|
||||
|
||||
def _get_flat_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads * head_size)
|
||||
|
||||
|
||||
# For testing sliced tensors
|
||||
def _get_padded_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads, head_size + 64)
|
||||
|
||||
|
||||
def _get_batch_tensor_shape(
|
||||
batch_size: int, seq_len: int, num_heads: int, head_size: int
|
||||
) -> tuple[int, ...]:
|
||||
return (batch_size, seq_len, num_heads, head_size)
|
||||
|
||||
|
||||
TENSORS_SHAPES_FN = [
|
||||
_get_batch_tensor_shape,
|
||||
_get_flat_tensor_shape,
|
||||
_get_padded_tensor_shape,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
|
||||
@pytest.mark.parametrize("tensor_shape_fn", TENSORS_SHAPES_FN)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_SIZES)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("use_key", USE_KEY)
|
||||
@torch.inference_mode()
|
||||
def test_rotary_embedding(
|
||||
default_vllm_config,
|
||||
is_neox_style: bool,
|
||||
tensor_shape_fn: Callable[[int, int, int, int], tuple[int, ...]],
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
rotary_dim: int | None,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
use_key: bool,
|
||||
max_position: int = 8192,
|
||||
rope_theta: float = 10000,
|
||||
) -> None:
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters = {
|
||||
"rope_type": "default",
|
||||
"rope_theta": rope_theta,
|
||||
"partial_rotary_factor": rotary_dim / head_size,
|
||||
}
|
||||
rope = get_rope(head_size, max_position, is_neox_style, rope_parameters)
|
||||
rope = rope.to(dtype=dtype, device=torch.get_default_device())
|
||||
|
||||
positions = torch.randint(0, max_position, (batch_size, seq_len))
|
||||
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
|
||||
# slice tensor if required, noop otherwise
|
||||
query = torch.randn(query_shape, dtype=dtype)[..., :head_size]
|
||||
key = torch.randn_like(query)[..., :head_size] if use_key else None
|
||||
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_query, ref_key = rope.forward_native(positions, query, key)
|
||||
out_query, out_key = rope.forward(positions, query, key)
|
||||
# Compare the results.
|
||||
torch.testing.assert_close(
|
||||
out_query,
|
||||
ref_query,
|
||||
atol=get_default_atol(out_query),
|
||||
rtol=get_default_rtol(out_query),
|
||||
)
|
||||
if use_key:
|
||||
torch.testing.assert_close(
|
||||
out_key,
|
||||
ref_key,
|
||||
atol=get_default_atol(out_key),
|
||||
rtol=get_default_rtol(out_key),
|
||||
)
|
||||
else:
|
||||
assert ref_key is None and out_key is None, "expected returned key to be None"
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_rope_module_cache(default_vllm_config):
|
||||
MAX_POSITIONS = [123, 1234]
|
||||
ROPE_THETAS = [10000, 1000000]
|
||||
ROPE_PARAMETERS = (
|
||||
{"rope_type": "default"},
|
||||
{"rope_type": "linear", "factor": (1,)},
|
||||
{"rope_type": "dynamic", "factor": 1},
|
||||
)
|
||||
settings = (
|
||||
HEAD_SIZES,
|
||||
ROTARY_DIMS,
|
||||
MAX_POSITIONS,
|
||||
ROPE_THETAS,
|
||||
IS_NEOX_STYLE,
|
||||
ROPE_PARAMETERS,
|
||||
DTYPES,
|
||||
)
|
||||
rope_setting_id_map: dict[str, int] = {}
|
||||
for setting in product(*settings):
|
||||
(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position,
|
||||
rope_theta,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
) = setting
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters["rope_theta"] = rope_theta
|
||||
rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
|
||||
rope = get_rope(
|
||||
head_size,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
)
|
||||
# different settings cannot share the same rope module
|
||||
assert id(rope) not in rope_setting_id_map.values()
|
||||
assert all(x.dtype == dtype for x in rope.buffers())
|
||||
assert all(x.dtype == dtype for x in rope.parameters())
|
||||
rope_setting_id_map[str(setting)] = id(rope)
|
||||
|
||||
for setting in product(*settings):
|
||||
(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position,
|
||||
rope_theta,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
) = setting
|
||||
if rotary_dim is None:
|
||||
rotary_dim = head_size
|
||||
rope_parameters["rope_theta"] = rope_theta
|
||||
rope_parameters["partial_rotary_factor"] = rotary_dim / head_size
|
||||
rope = get_rope(
|
||||
head_size,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rope_parameters,
|
||||
dtype,
|
||||
)
|
||||
# check if cache take effect
|
||||
assert id(rope) == rope_setting_id_map[str(setting)]
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for miscellaneous utilities
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
|
||||
|
||||
def rotary_embedding_opcheck(
|
||||
rot,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
):
|
||||
cos_sin_cache = rot.cos_sin_cache.to(query.device, dtype=query.dtype)
|
||||
|
||||
# ops.rotary_embedding() is a in-place operation
|
||||
# that updates the query and key tensors.
|
||||
opcheck(
|
||||
torch.ops._C.rotary_embedding,
|
||||
(positions, query, key, rot.head_size, cos_sin_cache, rot.is_neox_style),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cuda"])
|
||||
@pytest.mark.parametrize("max_position", [11, 4096, 32768])
|
||||
@pytest.mark.parametrize("is_neox_style", [True, False])
|
||||
@pytest.mark.parametrize("rotary_dim", [32])
|
||||
@pytest.mark.parametrize("head_size", [32, 108])
|
||||
@pytest.mark.parametrize("seq_len", [11, 1024])
|
||||
@pytest.mark.parametrize("use_key", [True, False])
|
||||
@pytest.mark.parametrize("head_stride_is_contiguous", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_rotary_embedding_opcheck(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
device,
|
||||
max_position,
|
||||
is_neox_style,
|
||||
rotary_dim,
|
||||
head_size,
|
||||
seq_len,
|
||||
use_key,
|
||||
head_stride_is_contiguous,
|
||||
dtype,
|
||||
):
|
||||
batch_size = 1
|
||||
base = 10000
|
||||
num_heads = 7
|
||||
rot = RotaryEmbedding(
|
||||
head_size, rotary_dim, max_position, base, is_neox_style, dtype
|
||||
)
|
||||
|
||||
positions = torch.randint(0, max_position, (batch_size, seq_len), device=device)
|
||||
head_stride = head_size + (64 if head_stride_is_contiguous else 0)
|
||||
|
||||
query = torch.randn(
|
||||
batch_size, seq_len, num_heads, head_stride, dtype=dtype, device=device
|
||||
)
|
||||
key = torch.randn_like(query) if use_key else None
|
||||
query = query[..., :head_size]
|
||||
key = key[..., :head_size] if key is not None else None
|
||||
|
||||
rotary_embedding_opcheck(rot, positions, query, key)
|
||||
|
||||
# if we have a contiguous head stride, test the alternate
|
||||
# [..., num_heads * head_dim] shape/layout
|
||||
if head_stride_is_contiguous:
|
||||
rotary_embedding_opcheck(
|
||||
rot,
|
||||
positions,
|
||||
query.flatten(start_dim=-2),
|
||||
key.flatten(start_dim=-2) if key is not None else None,
|
||||
)
|
||||
@@ -0,0 +1,210 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for fused MLA KV-cache write and RoPE fused kernel
|
||||
"""
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
|
||||
from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def default_vllm_config(monkeypatch):
|
||||
"""Enable the AITER triton rope on ROCm for fp16-consistent numerics.
|
||||
|
||||
The fused CUDA kernel runs native fp16 while forward_native upcasts to
|
||||
fp32, so on ROCm we route through the AITER triton rope (+rotary_embedding)
|
||||
to match. Its env gates are cached at import, hence refresh_env_variables().
|
||||
"""
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
|
||||
|
||||
is_rocm = current_platform.is_rocm()
|
||||
if is_rocm:
|
||||
config = VllmConfig(
|
||||
compilation_config=CompilationConfig(custom_ops=["+rotary_embedding"])
|
||||
)
|
||||
else:
|
||||
config = VllmConfig()
|
||||
try:
|
||||
with monkeypatch.context() as m, set_current_vllm_config(config):
|
||||
if is_rocm:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
m.setenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "1")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
yield config
|
||||
finally:
|
||||
if is_rocm:
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("is_neox_style", [False, True])
|
||||
@pytest.mark.parametrize("seq_len", [11, 42])
|
||||
@pytest.mark.parametrize("qk_rope_head_dim", [64, 128])
|
||||
@pytest.mark.parametrize("num_q_heads", [128])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
@pytest.mark.parametrize("kv_lora_rank", [512])
|
||||
@pytest.mark.parametrize("num_blocks", [64])
|
||||
@pytest.mark.parametrize("block_size", [16, 64, 256])
|
||||
@pytest.mark.parametrize("seed", [0])
|
||||
@pytest.mark.parametrize(
|
||||
"device",
|
||||
[f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)],
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_concat_and_cache_mla_rope_fused(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
is_neox_style: bool,
|
||||
seq_len: int,
|
||||
qk_rope_head_dim: int,
|
||||
num_q_heads: int,
|
||||
kv_cache_dtype: str,
|
||||
kv_lora_rank: int,
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
max_position: int = 8192,
|
||||
base: float = 10000,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
rope = RotaryEmbedding(
|
||||
qk_rope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
max_position,
|
||||
base,
|
||||
is_neox_style,
|
||||
torch.float32,
|
||||
)
|
||||
|
||||
rope = rope.to(dtype=dtype, device=torch.get_default_device())
|
||||
|
||||
positions = torch.randint(0, max_position, (seq_len,))
|
||||
|
||||
query = torch.randn(seq_len, num_q_heads, qk_rope_head_dim, dtype=dtype)
|
||||
key = torch.randn(seq_len, 1, qk_rope_head_dim + kv_lora_rank, dtype=dtype)
|
||||
|
||||
k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
|
||||
kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# We use forward_hip for the same numerics as the fused custom kernel on ROCm
|
||||
# when dtype is FP16. The torch-native implementation implicitly upcasts
|
||||
# FP16 x FP16 multiplications to FP32 before downcasting them, which leads
|
||||
# to notable output divergences.
|
||||
# Clone the tensors because the implementation modifies them in-place
|
||||
ref_q_pe, ref_k_pe = rope.forward_hip(positions, query.clone(), k_pe.clone())
|
||||
else:
|
||||
# NOTE(woosuk): The reference implementation should be executed first
|
||||
# because the custom kernel is in-place.
|
||||
ref_q_pe, ref_k_pe = rope.forward_native(positions, query, k_pe)
|
||||
assert ref_k_pe is not None
|
||||
|
||||
ref_k_pe = torch.flatten(ref_k_pe, start_dim=1).to(device=device)
|
||||
ref_k_rope = ref_k_pe[..., :qk_rope_head_dim]
|
||||
|
||||
total_available_slots = num_blocks * block_size
|
||||
total_needed_slots = seq_len
|
||||
assert total_available_slots >= total_needed_slots, "Not enough kv slots!"
|
||||
|
||||
slot_mapping_lst = random.sample(range(total_available_slots), total_needed_slots)
|
||||
slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
|
||||
|
||||
entry_size = kv_lora_rank + qk_rope_head_dim
|
||||
|
||||
kv_cache_scale = torch.tensor([0.1], dtype=torch.float32, device=device)
|
||||
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
entry_size,
|
||||
dtype=torch.uint8 if kv_cache_dtype == "fp8" else dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)
|
||||
|
||||
for i in range(seq_len):
|
||||
slot = slot_mapping[i].item()
|
||||
block_idx = slot // block_size
|
||||
block_offset = slot % block_size
|
||||
ref_temp[block_idx, block_offset] = torch.cat((kv_c[i], ref_k_rope[i]), -1)
|
||||
|
||||
if kv_cache_dtype == "fp8":
|
||||
ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
|
||||
ops.convert_fp8(
|
||||
ref_kv_cache, ref_temp, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
|
||||
)
|
||||
else:
|
||||
ref_kv_cache = ref_temp
|
||||
|
||||
opcheck(
|
||||
torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused,
|
||||
(
|
||||
positions,
|
||||
query,
|
||||
k_pe,
|
||||
kv_c,
|
||||
rope.cos_sin_cache,
|
||||
is_neox_style,
|
||||
slot_mapping,
|
||||
kv_cache,
|
||||
kv_cache_dtype,
|
||||
kv_cache_scale,
|
||||
),
|
||||
test_utils=DEFAULT_OPCHECK_TEST_UTILS,
|
||||
)
|
||||
|
||||
ops.concat_and_cache_mla_rope_fused(
|
||||
positions,
|
||||
query,
|
||||
k_pe,
|
||||
kv_c,
|
||||
rope.cos_sin_cache,
|
||||
is_neox_style,
|
||||
slot_mapping,
|
||||
kv_cache,
|
||||
kv_cache_dtype,
|
||||
kv_cache_scale,
|
||||
)
|
||||
|
||||
# ROCm neox-style Triton FMA diverges slightly from the fused kernel, so
|
||||
# relax the affected tolerance: rtol for fp8 (one e4m3 ULP ~12.5%) and atol
|
||||
# otherwise (bounded ~6e-4). Other paths use the CUDA defaults.
|
||||
rocm_neox = current_platform.is_rocm() and is_neox_style
|
||||
if kv_cache_dtype == "fp8":
|
||||
result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
|
||||
ops.convert_fp8(
|
||||
result_temp,
|
||||
kv_cache.contiguous(),
|
||||
kv_cache_scale.item(),
|
||||
kv_dtype=kv_cache_dtype,
|
||||
)
|
||||
expected_temp = torch.empty_like(ref_kv_cache, dtype=torch.float16)
|
||||
ops.convert_fp8(
|
||||
expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
result_temp, expected_temp, atol=0.001, rtol=0.15 if rocm_neox else 0.1
|
||||
)
|
||||
elif rocm_neox:
|
||||
torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
|
||||
else:
|
||||
torch.testing.assert_close(kv_cache, ref_kv_cache)
|
||||
|
||||
torch.testing.assert_close(
|
||||
query, ref_q_pe, atol=get_default_atol(query), rtol=get_default_rtol(query)
|
||||
)
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.utils.platform_utils import is_uva_available
|
||||
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
|
||||
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_cpu_write(device):
|
||||
torch.set_default_device(device)
|
||||
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
|
||||
cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
|
||||
assert cuda_view.device.type == "cuda"
|
||||
|
||||
assert cuda_view[0, 0] == 0
|
||||
assert cuda_view[2, 3] == 0
|
||||
assert cuda_view[4, 5] == 0
|
||||
|
||||
cpu_tensor[0, 0] = 1
|
||||
cpu_tensor[2, 3] = 2
|
||||
cpu_tensor[4, 5] = -1
|
||||
|
||||
cuda_view.mul_(2)
|
||||
assert cuda_view[0, 0] == 2
|
||||
assert cuda_view[2, 3] == 4
|
||||
assert cuda_view[4, 5] == -2
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_gpu_write(device):
|
||||
torch.set_default_device(device)
|
||||
cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32)
|
||||
cuda_view = get_accelerator_view_from_cpu_tensor(cpu_tensor)
|
||||
assert cuda_view.device.type == "cuda"
|
||||
|
||||
assert cuda_view[0, 0] == 0
|
||||
assert cuda_view[2, 3] == 0
|
||||
assert cuda_view[4, 5] == 0
|
||||
|
||||
cuda_view[0, 0] = 1
|
||||
cuda_view[2, 3] = 2
|
||||
cuda_view[4, 5] = -1
|
||||
cuda_view.mul_(2)
|
||||
|
||||
assert cpu_tensor[0, 0] == 2
|
||||
assert cpu_tensor[2, 3] == 4
|
||||
assert cpu_tensor[4, 5] == -2
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Accuracy tests for the fused Triton bilinear position-embedding kernel.
|
||||
|
||||
Compares ``triton_pos_embed_interpolate`` against the pure-PyTorch
|
||||
``pos_embed_interpolate_native`` across a variety of grid shapes and dtypes.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.model_executor.models.qwen3_vl import (
|
||||
pos_embed_interpolate_native,
|
||||
triton_pos_embed_interpolate,
|
||||
)
|
||||
|
||||
|
||||
DTYPES = [torch.float32, torch.bfloat16]
|
||||
# Qwen3-VL default
|
||||
NUM_GRID_PER_SIDE = 48
|
||||
SPATIAL_MERGE_SIZE = 2
|
||||
HIDDEN_DIM = 1152
|
||||
|
||||
# 4 square + 4 non-square grids (h, w divisible by spatial_merge_size=2)
|
||||
SQUARE_GRIDS = [(1, 4, 4), (1, 16, 16), (1, 32, 32), (1, 48, 48)]
|
||||
NON_SQUARE_GRIDS = [(1, 8, 16), (1, 14, 20), (1, 32, 48), (1, 60, 80)]
|
||||
ALL_GRIDS = SQUARE_GRIDS + NON_SQUARE_GRIDS
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
|
||||
@pytest.mark.parametrize(
|
||||
"grid_thw",
|
||||
ALL_GRIDS,
|
||||
ids=[f"{t}x{h}x{w}" for t, h, w in ALL_GRIDS],
|
||||
)
|
||||
def test_triton_matches_native(
|
||||
grid_thw: tuple[int, int, int],
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Triton kernel output must match the native PyTorch implementation."""
|
||||
t, h, w = grid_thw
|
||||
device = "cuda"
|
||||
|
||||
# Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
|
||||
torch.manual_seed(42)
|
||||
embed_weight = (
|
||||
torch.randn(
|
||||
NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
|
||||
HIDDEN_DIM,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
* 0.25
|
||||
)
|
||||
|
||||
native_out = pos_embed_interpolate_native(
|
||||
embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
|
||||
)
|
||||
triton_out = triton_pos_embed_interpolate(
|
||||
embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
|
||||
)
|
||||
|
||||
assert native_out.shape == triton_out.shape, (
|
||||
f"Shape mismatch: native {native_out.shape} vs triton {triton_out.shape}"
|
||||
)
|
||||
|
||||
# Small numerical differences arise from the precomputed h/w_scale
|
||||
# in the triton kernel vs torch.linspace in the native path, which can
|
||||
# cause single-ULP output differences
|
||||
# in a handful of elements.
|
||||
atol = {torch.float32: 5e-5, torch.bfloat16: 1e-2}[dtype]
|
||||
rtol = {torch.float32: 1e-5, torch.bfloat16: 1e-2}[dtype]
|
||||
torch.testing.assert_close(triton_out, native_out, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
|
||||
def test_temporal_repeat(dtype: torch.dtype) -> None:
|
||||
"""Verify temporal dimension t > 1 correctly repeats the spatial pattern."""
|
||||
device = "cuda"
|
||||
h, w = 16, 16
|
||||
t_single, t_multi = 1, 3
|
||||
|
||||
# Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
|
||||
torch.manual_seed(42)
|
||||
embed_weight = (
|
||||
torch.randn(
|
||||
NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
|
||||
HIDDEN_DIM,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
* 0.25
|
||||
)
|
||||
|
||||
out_single = triton_pos_embed_interpolate(
|
||||
embed_weight,
|
||||
t_single,
|
||||
h,
|
||||
w,
|
||||
NUM_GRID_PER_SIDE,
|
||||
SPATIAL_MERGE_SIZE,
|
||||
dtype,
|
||||
)
|
||||
out_multi = triton_pos_embed_interpolate(
|
||||
embed_weight,
|
||||
t_multi,
|
||||
h,
|
||||
w,
|
||||
NUM_GRID_PER_SIDE,
|
||||
SPATIAL_MERGE_SIZE,
|
||||
dtype,
|
||||
)
|
||||
|
||||
expected = out_single.repeat(t_multi, 1)
|
||||
torch.testing.assert_close(out_multi, expected, atol=0, rtol=0)
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
|
||||
|
||||
import contextlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
|
||||
def _has_flashinfer_cudnn() -> bool:
|
||||
"""Check if FlashInfer cuDNN backend is available."""
|
||||
try:
|
||||
from flashinfer.prefill import (
|
||||
cudnn_batch_prefill_with_kv_cache, # noqa: F401
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
HEAD_DIMS = [72, 80]
|
||||
SEQ_LENS = [256]
|
||||
NUM_HEADS = [16]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _fp8_attention():
|
||||
"""Create FP8-enabled MMEncoderAttention via config."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
|
||||
|
||||
mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
# MMEncoderAttention reads torch.get_default_dtype() during init
|
||||
# to determine the output dtype. In real model loading this is bf16.
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _build_cu_seqlens_and_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8_padded_hidden_size: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
|
||||
import numpy as np
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
|
||||
|
||||
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
torch.device("cuda"),
|
||||
)
|
||||
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(
|
||||
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
|
||||
),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
num_heads * head_dim,
|
||||
1, # tp_size
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
return cu_seqlens, max_seqlen, sequence_lengths
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_attn_output_shape(
|
||||
head_dim: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
_fp8_attention,
|
||||
) -> None:
|
||||
"""Verify FP8 attention produces correct output shape after un-padding."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
attn = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn is not None # mypy narrowing
|
||||
|
||||
# FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
|
||||
cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
|
||||
)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
|
||||
|
||||
# Output should have original head_dim (un-padded)
|
||||
assert output.shape[-1] == head_dim
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_vs_bf16_close(
|
||||
head_dim: int, seq_len: int, num_heads: int, _fp8_attention
|
||||
) -> None:
|
||||
"""FP8 attention output should be reasonably close to BF16 baseline."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
1,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
# FP8 path
|
||||
attn_fp8 = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn_fp8 is None or not attn_fp8.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn_fp8 is not None # mypy narrowing
|
||||
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
out_fp8 = attn_fp8._forward_flashinfer(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
seq_lengths,
|
||||
)
|
||||
|
||||
# BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
|
||||
# and calling the wrapper directly without FP8 quantization)
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.ops.vit_attn_wrappers import (
|
||||
vit_flashinfer_wrapper,
|
||||
)
|
||||
|
||||
out_bf16 = vit_flashinfer_wrapper(
|
||||
q=q.clone(),
|
||||
k=k.clone(),
|
||||
v=v.clone(),
|
||||
scale=attn_fp8.scale,
|
||||
workspace_buffer=_get_flashinfer_workspace_buffer(),
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
out_fp8_f = out_fp8.float()
|
||||
out_bf16_f = out_bf16.float()
|
||||
|
||||
abs_diff = (out_fp8_f - out_bf16_f).abs()
|
||||
abs_diff_flat = abs_diff.flatten()
|
||||
|
||||
# Relative diff (avoid division by zero)
|
||||
denom = out_bf16_f.abs().clamp(min=1e-6)
|
||||
rel_diff_flat = (abs_diff / denom).flatten()
|
||||
|
||||
cosine_sim = torch.nn.functional.cosine_similarity(
|
||||
out_fp8_f.flatten().unsqueeze(0),
|
||||
out_bf16_f.flatten().unsqueeze(0),
|
||||
).item()
|
||||
|
||||
pcts = [50, 90, 95, 99, 99.9]
|
||||
abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
|
||||
rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
|
||||
|
||||
print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
|
||||
print(f" cosine_sim={cosine_sim:.6f}")
|
||||
print(
|
||||
f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
|
||||
f"mean={abs_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
print(
|
||||
f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
|
||||
f"mean={rel_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
|
||||
assert abs_diff_flat.max().item() < 0.3, (
|
||||
f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
|
||||
)
|
||||
assert abs_diff_flat.mean().item() < 0.03, (
|
||||
f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
|
||||
)
|
||||
assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.kernels.triton.qkv_padded_fp8_quant import (
|
||||
quantize_fp8_pad_head_dim_triton,
|
||||
)
|
||||
|
||||
HEAD_DIMS = [72, 80, 128]
|
||||
SEQ_LENS = [64, 256]
|
||||
NUM_HEADS = [16]
|
||||
SCALES = [0.01, 0.1, 1.0]
|
||||
|
||||
|
||||
def _naive_fp8_quantize(
|
||||
tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
|
||||
) -> torch.Tensor:
|
||||
"""Reference FP8 quantization in PyTorch."""
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
|
||||
x = tensor.float()
|
||||
if not skip_scale:
|
||||
x = x / scale.item()
|
||||
x = x.clamp(fp8_min, fp8_max)
|
||||
return x.to(fp8_dtype)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("scale_val", SCALES)
|
||||
def test_quantize_contiguous(
|
||||
head_dim: int, seq_len: int, num_heads: int, scale_val: float
|
||||
) -> None:
|
||||
"""Test quantization of contiguous 3D tensors."""
|
||||
torch.manual_seed(42)
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
|
||||
1, 1, 1, 1
|
||||
)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
assert result.dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Compare unpadded portion against reference
|
||||
ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
# Padded region should be zero
|
||||
if padded_dim > head_dim:
|
||||
assert (result[:, :, head_dim:].float() == 0).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", [72, 80])
|
||||
def test_quantize_non_contiguous(head_dim: int) -> None:
|
||||
"""Test quantization from non-contiguous QKV views (interleaved buffer)."""
|
||||
seq_len, num_heads = 64, 16
|
||||
# Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
|
||||
qkv = torch.randn(
|
||||
seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
# Q is every 3rd head slice - non-contiguous view
|
||||
q = qkv[:, 0::3, :]
|
||||
assert not q.is_contiguous()
|
||||
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
result = quantize_fp8_pad_head_dim_triton(q, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
|
||||
# Compare against contiguous reference
|
||||
ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_skip_scale() -> None:
|
||||
"""Test skip_scale=True produces cast-only output (no division)."""
|
||||
seq_len, num_heads, head_dim = 32, 8, 80
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
|
||||
result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
|
||||
|
||||
# skip_scale should just cast, not divide
|
||||
ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
|
||||
torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
|
||||
|
||||
# With scale != 1.0, skip and no-skip should differ
|
||||
assert not torch.equal(result_skip.float(), result_noskip.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_4d_input() -> None:
|
||||
"""Test that 4D input (B, S, H, D) is handled correctly."""
|
||||
B, S, H, D = 2, 32, 8, 72
|
||||
tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
padded_dim = (D + 15) // 16 * 16
|
||||
assert result.shape == (B, S, H, padded_dim)
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_FP8_AMAX_HISTORY_LEN,
|
||||
_FP8_MAX,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
|
||||
LAYER_0 = "visual.blocks.0.attn.attn"
|
||||
LAYER_1 = "visual.blocks.1.attn.attn"
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _build_attention(mm_config):
|
||||
"""Yield an MMEncoderAttention with the given multimodal config.
|
||||
|
||||
The VllmConfig context stays active while the test runs so that
|
||||
``get_multimodal_config()`` calls during the forward path resolve. Also
|
||||
invokes ``process_weights_after_loading`` to simulate the model loader's
|
||||
auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
|
||||
"""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
yield None
|
||||
return
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_attention():
|
||||
"""Create an MMEncoderAttention with dynamic FP8 scaling."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_static_attention(tmp_path):
|
||||
"""Create an MMEncoderAttention with static FP8 scales from a file."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
|
||||
}
|
||||
)
|
||||
)
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
def test_dynamic_scaling_updates_scales(_make_attention) -> None:
|
||||
"""Verify that _record_amax_and_update_scales updates scale buffers."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
|
||||
S, H, D = 32, NUM_HEADS, HEAD_DIM
|
||||
q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
expected_q_scale = 2.0 / _FP8_MAX
|
||||
expected_k_scale = 3.0 / _FP8_MAX
|
||||
expected_v_scale = 4.0 / _FP8_MAX
|
||||
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
|
||||
torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
|
||||
torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
|
||||
|
||||
|
||||
def test_circular_buffer_wraps(_make_attention) -> None:
|
||||
"""Verify the amax circular buffer wraps at HISTORY_LEN."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN + 2):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
assert attn._fp8_amax_pos == 2
|
||||
|
||||
expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
|
||||
expected_scale = expected_max / _FP8_MAX
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
|
||||
|
||||
|
||||
def test_static_scales_loaded(_make_static_attention) -> None:
|
||||
"""Verify static scales are loaded from the JSON file."""
|
||||
attn = _make_static_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
assert attn.fp8_enabled
|
||||
assert not attn._fp8_dynamic_scale
|
||||
|
||||
# Layer 0 scales (the layer this attention was created with).
|
||||
assert attn._fp8_q_scale.item() == 224.0
|
||||
assert attn._fp8_k_scale.item() == 198.0
|
||||
assert attn._fp8_v_scale.item() == 210.0
|
||||
|
||||
assert not attn.skip_scale_q
|
||||
assert not attn.skip_scale_k
|
||||
assert not attn.skip_scale_v
|
||||
|
||||
# No amax history buffers for static scaling.
|
||||
assert not hasattr(attn, "_fp8_q_amax")
|
||||
|
||||
|
||||
def test_static_scales_missing_layer(tmp_path) -> None:
|
||||
"""Verify error when requested layer is not in the scale file."""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
scale_file = tmp_path / "wrong_layer.json"
|
||||
scale_file.write_text(
|
||||
json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
|
||||
)
|
||||
mm_config = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="scales not found for layer"):
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
|
||||
|
||||
def test_dynamic_scales_auto_save(tmp_path) -> None:
|
||||
"""Verify scales are saved to disk after the amax buffer fills."""
|
||||
import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
# Reset module-level state between runs (other tests may have left
|
||||
# state behind after triggering a save).
|
||||
_mod._fp8_scale_save_path = None
|
||||
_mod._fp8_saved_scale_refs.clear()
|
||||
|
||||
save_file = tmp_path / "auto_scales.json"
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_save_path=str(save_file),
|
||||
)
|
||||
) as attn:
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
# Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
# File should have been written on the 16th call (buffer wrap).
|
||||
assert save_file.is_file(), "Scale file was not saved"
|
||||
scales = json.loads(save_file.read_text())
|
||||
assert LAYER_0 in scales
|
||||
assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
|
||||
for val in scales[LAYER_0].values():
|
||||
assert isinstance(val, float) and val > 0
|
||||
|
||||
# Path is cleared after the one-shot save fires.
|
||||
assert _mod._fp8_scale_save_path is None
|
||||
Reference in New Issue
Block a user