obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-17 11:10:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Model]Fix MiniMaxM2ForCausalLM perf regression (#45935)

Browse Source 
Signed-off-by: Jee Jee Li <[email protected]>
This commit is contained in:
Jee Jee Li
2026-06-22 00:28:52 +08:00
committed by GitHub
parent 2cac89f9da
commit 745bba5ea8
2 changed files with 223 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/core/test_minimax_reduce_rms.py
+59 -3
View File
@@ -10,8 +10,12 @@ 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
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
@@ -54,8 +58,19 @@ def _worker_forward_qk(
torch.manual_seed(seed + 1000 + local_rank)
qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
# 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
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
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)
vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py
+164 -21
View File
@@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import (
)
from vllm.logger import init_logger
from vllm.model_executor.custom_op import CustomOp
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, tl, triton
from vllm.utils.torch_utils import direct_register_custom_op
logger = init_logger(__name__)
@@ -40,8 +40,116 @@ def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
return tensor_model_parallel_all_reduce(var.flatten()).view_as(var)
@torch.compile(backend=current_platform.simple_compile_backend, dynamic=True)
def _minimax_qk_norm_fallback(
@triton.jit
def _minimax_qk_var_kernel(
qkv_ptr, # [num_tokens, hidden], 16-bit activations
var_ptr, # [num_tokens, 2], fp32
row_stride, # element stride between tokens in qkv
q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides
kv_size: tl.constexpr,
BLOCK: tl.constexpr,
):
"""TP-pre stage: per-token mean-of-squares for the q and k segments.
Accumulates in fp32 while reading the 16-bit qkv in place, so no fp32
copy of q/k is materialized. ``var[:, 0]`` is the q variance and
``var[:, 1]`` the k variance; both are the local-shard means, ready for
the all-reduce that follows.
"""
token = tl.program_id(0)
base = qkv_ptr + token * row_stride
q_acc = 0.0
for off in range(0, q_size, BLOCK):
idx = off + tl.arange(0, BLOCK)
mask = idx < q_size
x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
q_acc += tl.sum(x * x, axis=0)
k_acc = 0.0
for off in range(0, kv_size, BLOCK):
idx = off + tl.arange(0, BLOCK)
mask = idx < kv_size
x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
k_acc += tl.sum(x * x, axis=0)
tl.store(var_ptr + token * 2 + 0, q_acc / q_size)
tl.store(var_ptr + token * 2 + 1, k_acc / kv_size)
@triton.jit
def _minimax_rms_apply_kernel(
qkv_ptr, # [num_tokens, hidden]
var_ptr, # [num_tokens, 2], fp32, all-reduced sum of per-shard means
q_w_ptr, # [q_size], q per-channel weight
k_w_ptr, # [kv_size], k per-channel weight
q_out_ptr, # [num_tokens, q_size], contiguous
k_out_ptr, # [num_tokens, kv_size], contiguous
row_stride, # element stride between tokens in qkv
q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides
kv_size: tl.constexpr,
tp_world: tl.constexpr, # folds the post-all-reduce /tp_world into rsqrt
eps: tl.constexpr,
BLOCK: tl.constexpr,
):
"""TP-post stage: ``x * rsqrt(var / tp_world + eps) * weight``.
A single program normalizes both the q and k segments of one token, so q
and k share one launch instead of two. The all-reduce yields the sum of
per-shard means, so the ``/ tp_world`` that recovers the global
mean-of-squares is folded into the ``rsqrt`` here rather than run as a
separate elementwise pass over the ``[num_tokens, 2]`` variance tensor.
"""
token = tl.program_id(0)
base = qkv_ptr + token * row_stride
q_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 0) / tp_world + eps)
q_out_row = q_out_ptr + token * q_size
for off in range(0, q_size, BLOCK):
idx = off + tl.arange(0, BLOCK)
mask = idx < q_size
x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
w = tl.load(q_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
y = x * q_inv * w
tl.store(q_out_row + idx, y.to(q_out_ptr.dtype.element_ty), mask=mask)
k_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 1) / tp_world + eps)
k_out_row = k_out_ptr + token * kv_size
for off in range(0, kv_size, BLOCK):
idx = off + tl.arange(0, BLOCK)
mask = idx < kv_size
x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
w = tl.load(k_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
y = x * k_inv * w
tl.store(k_out_row + idx, y.to(k_out_ptr.dtype.element_ty), mask=mask)
def _minimax_qk_norm_tp_eager(
qkv: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
q_size: int,
kv_size: int,
tp_world: int,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Pure-torch reference path used when Triton is unavailable."""
q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
orig_dtype = q.dtype
q = q.to(torch.float32)
k = k.to(torch.float32)
q_var = q.pow(2).mean(dim=-1, keepdim=True)
k_var = k.pow(2).mean(dim=-1, keepdim=True)
qk_var = torch.cat([q_var, k_var], dim=-1)
qk_var = _all_reduce_variance(qk_var) / tp_world
q_var, k_var = qk_var.chunk(2, dim=-1)
q = q * torch.rsqrt(q_var + eps) * q_weight
k = k * torch.rsqrt(k_var + eps) * k_weight
return q.to(orig_dtype), k.to(orig_dtype)
def _minimax_qk_norm_tp_fallback(
qkv: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
@@ -51,19 +159,50 @@ def _minimax_qk_norm_fallback(
tp_world: int,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
orig_dtype = q.dtype
q = q.to(torch.float32)
k = k.to(torch.float32)
q_var = q.pow(2).mean(dim=-1, keepdim=True)
k_var = k.pow(2).mean(dim=-1, keepdim=True)
if tp_world > 1:
qk_var = torch.cat([q_var, k_var], dim=-1)
qk_var = _all_reduce_variance(qk_var) / tp_world
q_var, k_var = qk_var.chunk(2, dim=-1)
q = q * torch.rsqrt(q_var + eps) * q_weight
k = k * torch.rsqrt(k_var + eps) * k_weight
return q.to(orig_dtype), k.to(orig_dtype)
"""All-reduce + QK RMSNorm without the Lamport fused kernel.
The all-reduce is a TP communication barrier and cannot live inside a
single kernel, so the eager-torch path is split into two Triton kernels
around it: a variance reduction before the all-reduce and a normalize
after. Compared to the ``torch.compile`` path this avoids materializing
fp32 copies of q/k and the ``cat``/``chunk`` temporaries.
"""
if not HAS_TRITON:
return _minimax_qk_norm_tp_eager(
qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
)
num_tokens = qkv.shape[0]
row_stride = qkv.stride(0)
BLOCK = 1024
grid = (num_tokens,)
qk_var = torch.empty(num_tokens, 2, dtype=torch.float32, device=qkv.device)
_minimax_qk_var_kernel[grid](
qkv, qk_var, row_stride, q_size=q_size, kv_size=kv_size, BLOCK=BLOCK
)
# All-reduce sums the per-shard means; the /tp_world that turns this back
# into the global mean is folded into the apply kernel's rsqrt below.
qk_var = _all_reduce_variance(qk_var)
q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device)
k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device)
_minimax_rms_apply_kernel[grid](
qkv,
qk_var,
q_weight,
k_weight,
q_out,
k_out,
row_stride,
q_size=q_size,
kv_size=kv_size,
tp_world=tp_world,
eps=eps,
BLOCK=BLOCK,
)
return q_out, k_out
def _minimax_qk_norm_fusion(
@@ -96,7 +235,7 @@ def _minimax_qk_norm_fusion(
tp_world,
eps,
)
return _minimax_qk_norm_fallback(
return _minimax_qk_norm_tp_fallback(
qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
)
@@ -231,10 +370,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
k = k.to(torch.float32)
q_var = q.pow(2).mean(dim=-1, keepdim=True)
k_var = k.pow(2).mean(dim=-1, keepdim=True)
if q_norm.tp_world > 1:
qk_var = torch.cat([q_var, k_var], dim=-1)
qk_var = _all_reduce_variance(qk_var) / q_norm.tp_world
q_var, k_var = qk_var.chunk(2, dim=-1)
q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
q = q.to(orig_dtype)
@@ -250,7 +386,14 @@ class MiniMaxText01RMSNormTP(CustomOp):
kv_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
assert qkv.ndim == 2
assert q_norm.variance_epsilon == k_norm.variance_epsilon
# Case 0： tp_size=1
if get_tensor_model_parallel_world_size() == 1:
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
q, k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q, k)
return q, k, v
# Case : tp_size>1
q, k = torch.ops.vllm.minimax_qk_norm_fusion(
qkv,
q_norm.weight,