diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py index de9fc2bbb4f..b9c591bb93f 100644 --- a/tests/kernels/core/test_minimax_reduce_rms.py +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -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) diff --git a/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py b/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py index e2c938ddad1..e48d9c01354 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py +++ b/vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py @@ -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,