Fix MXFP4 mlp_forward to handle 2D and 3D hidden_states shapes for multi-turn chat

src/transformers/integrations/mxfp4.py

@@ -289,16 +289,21 @@ def mlp_forward(self, hidden_states):
     else:
         routing = triton_kernels_hub.routing.routing
 
-        routing = routing
-    batch_size = hidden_states.shape[0]
-    hidden_states = hidden_states.reshape(-1, self.router.hidden_dim)
+    is_3d = hidden_states.ndim == 3
+    if is_3d:
+        batch_size, seq_len, _ = hidden_states.shape
+        hidden_states = hidden_states.reshape(-1, self.router.hidden_dim)
+
     router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias)
 
     with torch.cuda.device(router_logits.device):
         routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k)
 
     routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
-    routed_out = routed_out.reshape(batch_size, -1, self.router.hidden_dim)
+
+    if is_3d:
+        routed_out = routed_out.reshape(batch_size, seq_len, self.router.hidden_dim)
+
     return routed_out, router_logits
 
 

tests/quantization/mxfp4/test_mxfp4.py

@@ -31,8 +31,13 @@
 
 
 if is_torch_available():
+    from types import SimpleNamespace
+
     import torch
 
+    from transformers.integrations.mxfp4 import mlp_forward
+    from transformers.testing_utils import torch_device
+
 
 class Mxfp4ConfigTest(unittest.TestCase):
     def test_basic_config_creation(self):
@@ -194,7 +199,7 @@ def test_quantizer_validation_missing_triton(self):
         """Test quantizer validation when triton is not available"""
         with (
             patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False),
-            patch("transformers.quantizers.quantizer_mxfp4.is_kernels_availalble", return_value=False),
+            patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False),
         ):
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
@@ -208,7 +213,7 @@ def test_quantizer_validation_missing_triton_pre_quantized_no_dequantize(self):
         """Test quantizer validation when triton is not available but model is pre-quantized and dequantize is False"""
         with (
             patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False),
-            patch("transformers.quantizers.quantizer_mxfp4.is_kernels_availalble", return_value=False),
+            patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False),
         ):
             from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
 
@@ -365,6 +370,63 @@ def test_quantize_to_mxfp4(self):
         self.assertIsNotNone(flex_data)
         self.assertIsNotNone(mx_ctx)
 
+    @require_torch
+    @patch("torch.cuda.device")
+    @patch("transformers.integrations.mxfp4.triton_kernels_hub", create=True)
+    def test_mlp_forward_dimensionality(self, mock_triton_hub, mock_cuda_device):
+        """
+        Tests that the `mlp_forward` function correctly handles 2D and 3D tensors,
+        preserving the input dimensionality in the output. This is a regression
+        test for an issue seen in multi-turn chat scenarios.
+        """
+        # 1. Setup
+        hidden_dim = 32
+        num_experts = 4
+        top_k = 1
+
+        # Mock the routing function from the triton kernel
+        mock_routing_data = SimpleNamespace()  # Dummy object
+        mock_gather_idx = torch.zeros(10, dtype=torch.int32, device=torch_device)
+        mock_scatter_idx = torch.zeros(10, dtype=torch.int32, device=torch_device)
+        mock_triton_hub.routing.routing.return_value = (mock_routing_data, mock_gather_idx, mock_scatter_idx)
+
+        # Mock the 'self' object that mlp_forward is bound to
+        mock_self = SimpleNamespace()
+        mock_self.router = SimpleNamespace(
+            hidden_dim=hidden_dim,
+            weight=torch.randn((num_experts * top_k, hidden_dim), device=torch_device),
+            bias=torch.randn(num_experts * top_k, device=torch_device),
+            top_k=top_k,
+        )
+
+        # The experts function should return a tensor of the same shape as its input
+        def mock_experts_forward(hidden_states, *args, **kwargs):
+            return hidden_states
+
+        mock_self.experts = mock_experts_forward
+
+        # 2. Test with 2D input
+        hidden_states_2d = torch.randn(10, hidden_dim, device=torch_device)
+        output_2d, _ = mlp_forward(mock_self, hidden_states_2d)
+
+        self.assertEqual(len(output_2d.shape), 2, "Output should be 2D for 2D input")
+        self.assertEqual(hidden_states_2d.shape, output_2d.shape, "Shape should be preserved for 2D input")
+
+        # 3. Test with 3D input
+        hidden_states_3d = torch.randn(2, 5, hidden_dim, device=torch_device)
+        # The mock routing indices need to match the flattened input size
+        num_tokens = hidden_states_3d.shape[0] * hidden_states_3d.shape[1]
+        mock_triton_hub.routing.routing.return_value = (
+            mock_routing_data,
+            torch.zeros(num_tokens, dtype=torch.int32, device=torch_device),
+            torch.zeros(num_tokens, dtype=torch.int32, device=torch_device),
+        )
+
+        output_3d, _ = mlp_forward(mock_self, hidden_states_3d)
+
+        self.assertEqual(len(output_3d.shape), 3, "Output should be 3D for 3D input")
+        self.assertEqual(hidden_states_3d.shape, output_3d.shape, "Shape should be preserved for 3D input")
+
 
 @require_torch
 @require_torch_large_gpu
@@ -479,3 +541,44 @@ def test_memory_footprint_comparison(self):
         quantized_mem = quantized_model.get_memory_footprint()
         dequantized_mem = dequantized_model.get_memory_footprint()
         self.assertLess(quantized_mem, dequantized_mem)
+
+    def test_gpt_oss_multi_turn_generation(self):
+        """
+        An integration test to ensure multi-turn chat works with GPT-OSS, specifically testing the fix
+        for the tensor shape mismatch in the MXFP4 MLP layer. This test was failing before the fix
+        when Mxfp4Config(dequantize=False) i.e using this quantized version
+        """
+
+        # we want dequantize=False to test for multi-turn
+        quantization_config = Mxfp4Config(dequantize=False)
+        model = GptOssForCausalLM.from_pretrained(
+            self.model_name,
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+
+        # Turn 1: Initial question
+        prompt_1 = "Who directed the movie Inception?"
+        inputs_1 = tokenizer(prompt_1, return_tensors="pt").to(model.device)
+
+        # This first generation uses a 3D tensor internally
+        output_1_ids = model.generate(**inputs_1, max_new_tokens=20)
+        response_1 = tokenizer.decode(output_1_ids[0], skip_special_tokens=True)
+
+        # Turn 2: Follow-up question
+        prompt_2 = response_1 + "\nAnd what other movies has he directed?"
+        inputs_2 = tokenizer(prompt_2, return_tensors="pt").to(model.device)
+
+        # This is the step that would previously crash with the `num_stages` assertion error.
+        try:
+            output_2_ids = model.generate(**inputs_2, max_new_tokens=50)
+            response_2 = tokenizer.decode(output_2_ids[0], skip_special_tokens=True)
+        except Exception as e:
+            self.fail(f"Multi-turn generation failed with MXFP4. Error: {e}")
+
+        # Assert that we got a coherent, non-empty response
+        # It will tell that our multi-turn has success
+        self.assertIn("Inception", response_1)
+        self.assertTrue(len(response_2) > len(prompt_2))