diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 53d760c339..d17416de99 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( FlashInferExperts, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + view_as_block_major_k, +) from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( TrtLlmFp8ExpertsModular, TrtLlmFp8ExpertsMonolithic, @@ -570,6 +573,10 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( padded = w2_shape[-1] w13_ptr = layer.w13_weight.data_ptr() w2_ptr = layer.w2_weight.data_ptr() + load_numel = ( + layer.w13_weight.weight_loader_numel, + layer.w2_weight.weight_loader_numel, + ) for _ in range(2): reloaded_w13 = torch.randn_like(layer.w13_weight) @@ -602,19 +609,24 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( assert layer.w2_weight.shape == w2_shape assert layer.w13_weight.data_ptr() == w13_ptr assert layer.w2_weight.data_ptr() == w2_ptr - kernel_w13, kernel_w2 = method._kernel_weights(layer) + assert ( + layer.w13_weight.weight_loader_numel, + layer.w2_weight.weight_loader_numel, + ) == load_numel + kernel_w13 = view_as_block_major_k(layer.w13_weight) + kernel_w2 = view_as_block_major_k(layer.w2_weight) assert torch.equal(kernel_w13, expected_w13) assert torch.equal(kernel_w2, expected_w2) -def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index): +def _check_flashinfer_ipc_weights(entries, expected, mode, device_index): from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader from vllm.model_executor.utils import weights_already_processed torch.accelerator.set_device_index(device_index) with pytest.MonkeyPatch.context() as monkeypatch: method, layer = _make_unquantized_flashinfer_test_layer( - monkeypatch, 192, is_gated, device="meta" + monkeypatch, 192, is_gated=True, device="meta" ) loader = object.__new__(IpcModelLoader) loader.mode = mode @@ -625,7 +637,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde with weights_already_processed(): method.process_weights_after_loading(layer) assert method.moe_kernel is not None - actual = method._kernel_weights(layer) + actual = ( + view_as_block_major_k(layer.w13_weight), + view_as_block_major_k(layer.w2_weight), + ) for weight, reference, pointer in zip(actual, expected, pointers): assert weight.data_ptr() == pointer assert torch.equal(weight.cpu(), reference) @@ -636,11 +651,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde torch.accelerator.synchronize() -@pytest.mark.parametrize("is_gated", [True, False]) @pytest.mark.parametrize("cache_ndim", [3, 4]) @pytest.mark.parametrize("mode", ["copy", "zero_copy"]) def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( - monkeypatch, is_gated, cache_ndim, mode + monkeypatch, cache_ndim, mode ): """A spawned IPC consumer restores views without transient method state.""" from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( @@ -649,13 +663,12 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry method, layer = _make_unquantized_flashinfer_test_layer( - monkeypatch, 192, is_gated, device="meta" + monkeypatch, 192, is_gated=True, device="meta" ) packed = convert_moe_weights_to_flashinfer_trtllm_block_layout( {}, torch.randn_like(layer.w13_weight, device="cuda"), torch.randn_like(layer.w2_weight, device="cuda"), - is_gated_act_gemm=is_gated, ) expected = [weight.cpu() for weight in packed] entries = {} @@ -667,7 +680,6 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( args=( entries, expected, - is_gated, mode, torch.accelerator.current_device_index(), ), diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index 423f181132..ef31de06d1 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -12,7 +12,6 @@ correctly handles this mismatch. import pytest import torch -from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, ) @@ -465,30 +464,6 @@ class TestUnquantizedTrtLlmPrePadding: assert hidden == 64 assert intermediate == expected - @pytest.mark.parametrize("is_act_and_mul", [True, False]) - def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool): - method = self._make_method(1344) - if not is_act_and_mul: - method.moe.activation = MoEActivation.RELU2_NO_MUL - method.moe.intermediate_size_per_partition = 1408 - layer = torch.nn.Module() - - method.create_weights( - layer=layer, - num_experts=2, - hidden_size=64, - intermediate_size_per_partition=1408, - params_dtype=torch.bfloat16, - ) - - up_mult = 2 if is_act_and_mul else 1 - assert layer.w13_weight.shape == (2, up_mult * 1408, 64) - assert layer.w2_weight.shape == (2, 64, 1408) - assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0 - if is_act_and_mul: - assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0 - assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0 - class TestLoadWeightsExpertBias: """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py index f7466c63c8..5e569e153a 100644 --- a/tests/kernels/moe/test_trtllm_bf16_moe.py +++ b/tests/kernels/moe/test_trtllm_bf16_moe.py @@ -12,6 +12,7 @@ import pytest import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.moe.utils import make_dummy_moe_config from tests.kernels.utils import torch_moe from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -27,6 +28,9 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( TrtLlmBf16ExpertsModular, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import ( + TrtLlmBf16LoRAExperts, +) from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, convert_to_unquantized_kernel_format, @@ -150,3 +154,48 @@ def test_trtllm_bf16_moe_modular_no_graph( atol=1e-1, rtol=2e-1, ) + + +@pytest.mark.parametrize("has_lora_delta", [False, True]) +@torch.inference_mode() +def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta): + """LoRA dispatch accepts both 3D parameters and legacy 4D packed views.""" + config = make_dummy_moe_config( + num_experts=128, hidden_dim=256, intermediate_size=128 + ) + experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG) + x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10 + w1, w2 = convert_to_unquantized_kernel_format( + UnquantizedMoeBackend.FLASHINFER_TRTLLM, + config, + torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10, + torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10, + ) + assert w1.ndim == w2.ndim == 3 + topk = ( + torch.tensor([[0], [1]], device="cuda", dtype=torch.int32), + torch.ones(2, 1, device="cuda", dtype=torch.float32), + ) + delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16) + + def invoke(w1, w2): + result = experts.invoke_routed_moe( + hidden_states=x, + w1=w1, + w2=w2, + topk_ids_and_weights=topk, + gemm1_lora_delta=delta if has_lora_delta else None, + global_num_experts=128, + a1q_scale=None, + output=torch.empty_like(x), + ) + if has_lora_delta: + # Only compare rows belonging to real tokens, excluding padding. + indices = result[2].flatten().long() + return result[0][indices], result[3][indices] + return result + + actual = invoke(w1, w2) + expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64)) + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 446acf46dd..54227d3cd1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -31,6 +31,15 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +def view_as_block_major_k(weight: torch.Tensor) -> torch.Tensor: + """View packed storage, including legacy 4D IPC cache entries.""" + if weight.ndim == 4: + return weight + experts, rows, cols = weight.shape + block_k = 128 // weight.element_size() + return weight.view(experts, cols // block_k, rows, block_k) + + class TrtLlmBf16ExpertsBase: """BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and monolithic interfaces. @@ -200,8 +209,8 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular) result = flashinfer.fused_moe.trtllm_bf16_routed_moe( topk_ids=(topk_ids, topk_weights), hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), num_experts=global_num_experts, top_k=topk_ids.size(1), n_group=None, @@ -281,8 +290,8 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), num_experts=global_num_experts, top_k=self.topk, n_group=num_expert_group, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 4dc5076cf6..1c5698932a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -34,6 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRACon from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( LoRAExpertsMixin, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + view_as_block_major_k, +) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) @@ -527,8 +530,8 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase): ret = flashinfer.fused_moe.trtllm_bf16_routed_moe( topk_ids=topk_ids_and_weights, hidden_states=hidden_states, - gemm1_weights=w1, - gemm2_weights=w2, + gemm1_weights=view_as_block_major_k(w1), + gemm2_weights=view_as_block_major_k(w2), gemm1_lora_delta=gemm1_lora_delta, num_experts=global_num_experts, top_k=self.topk, diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index d1437933c3..bcb00eecbb 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey @@ -349,6 +350,17 @@ def select_unquantized_moe_backend( ) +def unquantized_round_up_hidden_size_and_intermediate_size( + backend: UnquantizedMoeBackend, + hidden_size: int, + intermediate_size: int, +) -> tuple[int, int]: + """Round up dimensions before allocation to satisfy the selected kernel.""" + if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + intermediate_size = round_up(intermediate_size, 128) + return hidden_size, intermediate_size + + def convert_to_unquantized_kernel_format( unquantized_backend: UnquantizedMoeBackend, moe_config: FusedMoEConfig, @@ -377,13 +389,26 @@ def convert_to_unquantized_kernel_format( ) moe_config.intermediate_size_per_partition = padded_intermediate + # Reloads only overwrite checkpoint slices. An earlier in-place + # permutation can leave nonzero values in the raw padding slots. + unpadded = moe_config.intermediate_size_per_partition_unpadded + assert unpadded is not None + if padded_intermediate > unpadded: + w13_weight[:, unpadded:padded_intermediate].zero_() + if is_act_and_mul: + w13_weight[:, padded_intermediate + unpadded :].zero_() + w2_weight[:, :, unpadded:].zero_() + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} - w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout( + convert_moe_weights_to_flashinfer_trtllm_block_layout( _cache_permute_indices, w13_weight, w2_weight, is_gated_act_gemm=is_act_and_mul, ) + # Keep checkpoint-shaped parameters for reload/IPC. The experts create + # BlockMajorK views at dispatch without changing the underlying storage. + return w13_weight, w2_weight if ( unquantized_backend == UnquantizedMoeBackend.TRITON diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index a1b2471f5d..e0835f538c 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( convert_to_unquantized_kernel_format, make_unquantized_moe_kernel, select_unquantized_moe_backend, + unquantized_round_up_hidden_size_and_intermediate_size, ) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -34,7 +35,6 @@ from vllm.model_executor.utils import ( set_weight_attrs, ) from vllm.platforms import current_platform -from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -72,9 +72,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ): unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded assert unpadded_intermediate is not None - has_intermediate_padding = ( - intermediate_size_per_partition > unpadded_intermediate - ) if self.moe.is_act_and_mul: w13_up_dim = 2 * intermediate_size_per_partition else: @@ -83,22 +80,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): assert unpadded_hidden is not None unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1) # Fused gate_up_proj (column parallel) - w13_data = torch.empty( - num_experts, - w13_up_dim, - hidden_size, - dtype=params_dtype, + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_up_dim, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, ) - if has_intermediate_padding: - # The checkpoint loader only writes the unpadded slices. Zero the - # tails now so no full-size padding copy is needed after loading. - w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_() - if self.moe.is_act_and_mul: - w13_data[ - :, - intermediate_size_per_partition + unpadded_intermediate :, - ].zero_() - w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden @@ -111,15 +101,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): set_weight_attrs(w13_bias, extra_weight_attrs) w13_bias.weight_loader_numel = num_experts * unpadded_up # down_proj (row parallel) - w2_data = torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=params_dtype, + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, ) - if has_intermediate_padding: - w2_data[:, :, unpadded_intermediate:].zero_() - w2_weight = torch.nn.Parameter(w2_data, requires_grad=False) layer.register_parameter("w2_weight", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) w2_weight.weight_loader_numel = ( @@ -147,11 +137,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): act_dtype, moe_parallel_config, ) - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - intermediate_size_per_partition = round_up( - intermediate_size_per_partition, 128 - ) - return hidden_size, intermediate_size_per_partition + return unquantized_round_up_hidden_size_and_intermediate_size( + self.unquantized_backend, hidden_size, intermediate_size_per_partition + ) def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor: # Pad the weight tensor. This is an optimization on ROCm platform, which @@ -177,21 +165,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): w13: torch.Tensor, w2: torch.Tensor, ) -> None: - w13_parameter_shape = w13.shape - w2_parameter_shape = w2.shape - - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - unpadded = self.moe.intermediate_size_per_partition_unpadded - assert unpadded is not None - intermediate = w2.shape[-1] - # Reloads overwrite only checkpoint slices; the previous in-place - # permutation can leave nonzero values in the raw padding slots. - if intermediate > unpadded: - w13[:, unpadded:intermediate].zero_() - if self.moe.is_act_and_mul: - w13[:, intermediate + unpadded :].zero_() - w2[:, :, unpadded:].zero_() - # Shuffle weights to runtime format. w13_new, w2_new = convert_to_unquantized_kernel_format( self.unquantized_backend, @@ -199,10 +172,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): w13_weight=w13, w2_weight=w2, ) - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - w13_new = w13_new.view(w13_parameter_shape) - w2_new = w2_new.view(w2_parameter_shape) - # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__; # On the first call we replace the parameter normally. On subsequent # calls (e.g. RL weight updates that re-trigger @@ -231,25 +200,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): assert self.moe_kernel is not None self.moe_kernel.fused_experts.process_weights_after_loading(layer) - def _kernel_weights( - self, layer: "RoutedExperts" - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM: - return layer.w13_weight, layer.w2_weight - - # IPC caches restore tensor shapes/storage, not method attributes. - # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements. - def block_view(weight: torch.Tensor) -> torch.Tensor: - if weight.ndim == 4: - return weight - experts, rows, cols = weight.shape - return weight.view(experts, cols // 64, rows, 64) - - return ( - block_view(layer.w13_weight), - block_view(layer.w2_weight), - ) - def _init_moe_kernel(self, layer: "RoutedExperts") -> None: """Build the MoE kernel from the layer's current (shuffled) weights.""" self.moe_quant_config = self.get_fused_moe_quant_config(layer) @@ -354,11 +304,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert self.moe_kernel is not None - w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply( hidden_states=x, - w1=w13, - w2=w2, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -396,11 +345,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None - w13, w2 = self._kernel_weights(layer) return self.moe_kernel.apply_monolithic( x, - w13, - w2, + layer.w13_weight, + layer.w2_weight, router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts,