diff --git a/csrc/activation_kernels.cu b/csrc/activation_kernels.cu index 9e28c3343b..0341236f77 100644 --- a/csrc/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -1059,33 +1059,6 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d] }); } -void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d] - torch::stable::Tensor& input, // [E, T, 2 * d] - const torch::stable::Tensor& expert_num_tokens, - double beta, double linear_beta) { - int num_experts = input.size(0); - int max_num_tokens = input.size(1); - int d = input.size(2) / 2; - if (num_experts == 0 || max_num_tokens == 0 || d == 0) { - return; - } - constexpr int block_size = 256; - dim3 grid((d + block_size - 1) / block_size, num_experts); - dim3 block(block_size); - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "masked_situ_and_mul_kernel", [&] { - vllm::masked_situ_and_mul_kernel - <<>>( - out.mutable_data_ptr(), - input.const_data_ptr(), - expert_num_tokens.const_data_ptr(), max_num_tokens, d, - (float)beta, (float)linear_beta); - }); -} - #define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \ BETA) \ vllm::masked_act_and_mul_kernel, ACT_FIRST, \ diff --git a/csrc/ops.h b/csrc/ops.h index 549dc84a1c..6e40dc7e54 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -579,10 +579,6 @@ void situ_and_mul_quant( int64_t group_size = 0, std::optional num_valid_tokens = std::nullopt, int64_t topk = 1); -void masked_situ_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& expert_num_tokens, - double beta = 1.0, double linear_beta = -1.0); void masked_moe_activation(torch::stable::Tensor& out, torch::stable::Tensor& input, const torch::stable::Tensor& valid_token_counts, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a92d9cc114..223cf5e101 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -640,9 +640,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " "float beta=1.0, float linear_beta=-1.0, int group_size=0, " "Tensor? num_valid_tokens=None, int topk=1) -> ()"); - ops.def( - "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " - "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); ops.def( "masked_moe_activation(Tensor! out, Tensor input, Tensor " "valid_token_counts, str activation, float clamp_limit=0.0, float " @@ -875,7 +872,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); - ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index 85ec4903ee..9303e23553 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -353,46 +353,6 @@ def test_silu_and_mul_with_clamp( ) -@pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) -@torch.inference_mode() -def test_masked_situ_and_mul( - default_vllm_config, - linear_beta: float, - dtype: torch.dtype, -) -> None: - """Masked SITU computes valid expert rows and preserves padded zeros.""" - device = CUDA_DEVICES[0] - num_experts, max_num_tokens, d = 4, 7, 512 - beta = 1.5 - input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device) - expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device) - output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device) - - torch.ops._C.masked_situ_and_mul( - output, input, expert_num_tokens, beta, linear_beta - ) - - gate, up = input.float().chunk(2, dim=-1) - expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) - if linear_beta > 0: - up = linear_beta * torch.tanh(up / linear_beta) - expected = (expected * up).to(dtype) - for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()): - torch.testing.assert_close( - output[expert, :num_tokens], - expected[expert, :num_tokens], - atol=get_default_atol(output), - rtol=get_default_rtol(output), - ) - assert torch.count_nonzero(output[expert, num_tokens:]) == 0 - - opcheck( - torch.ops._C.masked_situ_and_mul, - (output, input, expert_num_tokens, beta, linear_beta), - ) - - MOE_ACTIVATION_CASES = [ pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"), pytest.param( diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 4afb45f503..7969c82dda 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -10,7 +10,6 @@ import torch.nn.functional as F import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( - _torch_hc_prenorm_gemm, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -1002,6 +1001,19 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm): torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) +def _torch_hc_prenorm_gemm( + x: torch.Tensor, + fn: torch.Tensor, + out: torch.Tensor, + sqrsum: torch.Tensor, +) -> None: + assert out.shape[0] == 1 + assert sqrsum.shape[0] == 1 + x_float = x.float() + out[0].copy_(x_float @ fn.t()) + sqrsum[0].copy_(x_float.square().sum(dim=-1)) + + @pytest.mark.skipif( not HAS_TILELANG_MHC, reason="TileLang MHC support required", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0f57c9ecb4..fadb053915 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -21,7 +21,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( ) if TYPE_CHECKING: - from aiter import ActivationType, QuantType + from aiter import ActivationType logger = init_logger(__name__) @@ -2202,38 +2202,6 @@ class rocm_aiter_ops: } return mapping.get(name) - @staticmethod - def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None": - """Given a quantization type as a string, returns the corresponding aiter QuantType enum. - Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128". - Returns None if the mapping fails. - - Args: - quant_type_str (str): Quantization type as string. - - Returns: - Aiter QuantType enum value, or None if not found. - - """ - try: - from aiter import QuantType - except ImportError: - return None - - if not isinstance(quant_type_str, str): - return None - - name = quant_type_str.strip().lower() - mapping = { - "no": QuantType.No, - "per_tensor": QuantType.per_Tensor, - "per_token": QuantType.per_Token, - "per_1x32": QuantType.per_1x32, - "per_1x128": QuantType.per_1x128, - "per_128x128": QuantType.per_128x128, - } - return mapping.get(name) - @classmethod @if_aiter_supported def is_enabled(cls) -> bool: diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 62741bbda1..2c0f4c12a0 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -6,19 +6,6 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import direct_register_custom_op -def _torch_hc_prenorm_gemm( - x: torch.Tensor, - fn: torch.Tensor, - out: torch.Tensor, - sqrsum: torch.Tensor, -) -> None: - assert out.shape[0] == 1 - assert sqrsum.shape[0] == 1 - x_float = x.float() - out[0].copy_(x_float @ fn.t()) - sqrsum[0].copy_(x_float.square().sum(dim=-1)) - - def _hc_prenorm_gemm_outputs( x: torch.Tensor, fn: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/expert_map_manager.py b/vllm/model_executor/layers/fused_moe/expert_map_manager.py index daa2b3ff8e..1035654378 100644 --- a/vllm/model_executor/layers/fused_moe/expert_map_manager.py +++ b/vllm/model_executor/layers/fused_moe/expert_map_manager.py @@ -359,13 +359,6 @@ class ExpertMapManager: return True return self._expert_map[global_id] != -1 - def get_local_expert_ids(self) -> list[int]: - """Get list of global IDs for experts on this rank.""" - if self._expert_map is None: - return list(range(self.global_num_experts)) - - return torch.where(self._expert_map != -1)[0].tolist() - def update( self, moe_parallel_config: FusedMoEParallelConfig, diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index 1180e2df36..eb3ff4fe27 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -12,35 +12,6 @@ from vllm.triton_utils import tl, triton logger = logging.getLogger(__name__) -def block_dequant( - x_q_block: torch.Tensor, - x_s: torch.Tensor, - block_size: list[int], -) -> torch.Tensor: - """This function conducts block-wise dequantization. - The inputs are block-wise quantization tensor `x_q_block`, - block-wise quantization scale and the block size. - The outputs are dequantized tensor. - """ - block_n, block_k = block_size[0], block_size[1] - n, k = x_q_block.shape - n_tiles = (n + block_n - 1) // block_n - k_tiles = (k + block_k - 1) // block_k - assert n_tiles == x_s.shape[0] - assert k_tiles == x_s.shape[1] - - x_dq_block = x_q_block.to(torch.float32) - - for i in range(k_tiles): - for j in range(n_tiles): - x_dq_block[ - j * block_n : min((j + 1) * block_n, n), - i * block_k : min((i + 1) * block_k, k), - ] *= x_s[j][i] - - return x_dq_block - - if current_platform.is_rocm(): @triton.jit diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index d38497b2ff..085f092341 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -483,16 +483,6 @@ class TopKTopPSampler(nn.Module): return xpu_sample(logits, k, p, self.logprobs_mode) -# Note: this is a workaround for -# https://github.com/pytorch/pytorch/pull/151218 -@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) -def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor: - probs = logits.softmax(dim=-1, dtype=torch.float32) - q = torch.empty_like(probs) - q.exponential_() - return probs.div(q).argmax(dim=-1).view(-1) - - def apply_top_k_top_p( logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None ) -> torch.Tensor: