[AMD][Quantization] Online MXFP4 quantization 4/N - NVFP4 to MXFP4 Online Requantization on AMD GPUs (#29328)

test/registered/quant/test_quark_mxfp4.py

@@ -6,6 +6,7 @@ import unittest
 from sglang.test.ci.ci_register import register_amd_ci
 
 register_amd_ci(est_time=106, suite="stage-b-test-1-gpu-small-amd-mi35x")
+import os
 import time
 from types import SimpleNamespace
 
@@ -87,20 +88,10 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
                 raise RuntimeError(f"Server {url} failed to start in {timeout}s")
             time.sleep(1)
 
-        # # Extract and display peak GPU memory from logs
-        combined_output = cls.stdout.getvalue() + cls.stderr.getvalue()
-
-        peak_memory_before_load = cls._extract_peak_memory_before_load(combined_output)
-        if is_cuda_alike() and not peak_memory_before_load:
-            raise ValueError("Should have found peak memory")
-        cls.peak_memory_before_load = float(peak_memory_before_load)
-
-        memory_increase_load_weights = cls._extract_memory_increase_load_weights(
-            combined_output
-        )
-        if is_cuda_alike() and not memory_increase_load_weights:
-            raise ValueError("Should have found memory increase in load_weights")
-        cls.memory_increase_load_weights = float(memory_increase_load_weights)
+        # Keep the raw server for memory numbers, which are parsed lazily by
+        # _test_peak_memory so subclasses that don't test memory (e.g. the
+        # NVFP4->MXFP4 accuracy-only class) don't require these log lines.
+        cls.combined_output = cls.stdout.getvalue() + cls.stderr.getvalue()
 
     @classmethod
     def _extract_peak_memory_before_load(cls, log_output):
@@ -115,8 +106,11 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
     @classmethod
     def _extract_memory_increase_load_weights(cls, log_output):
         """Extract memory increase during load_weights call."""
-        # Search for the log message pattern
-        pattern = r"Memory increase during load_weights:\s+([\d.]+)\s+GiB"
+        # Signed: the value is (free_before - free_after) around load_weights.
+        # When the on-device source representation is larger than the loaded
+        # result (e.g. requantizing to a more compact format), loading frees
+        # net memory and the reported increase is negative.
+        pattern = r"Memory increase during load_weights:\s+(-?[\d.]+)\s+GiB"
         match = re.search(pattern, log_output)
         if match:
             return match.group(1)
@@ -138,21 +132,33 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
         if not is_cuda_alike():
             self.skipTest("not is_cuda_alike")
 
+        peak_memory_before_load = self._extract_peak_memory_before_load(
+            self.combined_output
+        )
+        if not peak_memory_before_load:
+            raise ValueError("Should have found peak memory")
+        peak_memory_before_load = float(peak_memory_before_load)
+
+        memory_increase_load_weights = self._extract_memory_increase_load_weights(
+            self.combined_output
+        )
+        if not memory_increase_load_weights:
+            raise ValueError("Should have found memory increase in load_weights")
+        memory_increase_load_weights = float(memory_increase_load_weights)
+
         # NOTE: We can not simply rely on peak memory after `load_weights` as functions used
-        # in-between (e.g. FP8->MXFP4 requantization) during weight loading may have a higher peak memory footprint
+        # in-between (e.g. NVFP4->MXFP4 requantization) during weight loading may have a higher peak memory footprint
         # than simply the allocated weights.
         if add_peak_memory_before_load:
-            reference_gib = (
-                self.memory_increase_load_weights + self.peak_memory_before_load
-            )
+            reference_gib = memory_increase_load_weights + peak_memory_before_load
         else:
-            reference_gib = self.memory_increase_load_weights
+            reference_gib = memory_increase_load_weights
 
         assert reference_gib < threshold
 
         if test_start:
             # Weights initialized on meta device (not for dense BF16->MXFP4)
-            assert self.peak_memory_before_load < 5
+            assert peak_memory_before_load < 5
 
     def _test_gsm8k(self, accuracy_threshold):
         """Helper method to test GSM8K accuracy against a threshold."""
@@ -205,6 +211,39 @@ class TestOnlineQuantizationMemoryLoadMOE(TestOnlineQuantizationMemoryLoad):
         self._test_gsm8k(accuracy_threshold=0.89)
 
 
+class TestNVFP4ToMXFP4MOETP1(TestOnlineQuantizationMemoryLoad):
+    # ModelOpt NVFP4 export (quant_method="modelopt", quant_algo="NVFP4") =>
+    # Nvfp4SourceConfig(). Exercises the NVFP4 -> MXFP4 MoE requantization path:
+    # the per-expert dequantize_nvfp4 + dynamic_mxfp4_quant requant, and the w13
+    # gate/up weight_scale_2 split in _requantize_nvfp4_to_mxfp4.
+    model = "nvidia/Qwen3-30B-A3B-NVFP4"  # NVFP4 model
+    tp = 1
+
+    def test_gsm8k(self):
+        # Requantized NVFP4 -> MXFP4 observed accuracy: ~0.88
+        # (BF16 Qwen/Qwen3-30B-A3B reference: ~0.94).
+        self._test_gsm8k(accuracy_threshold=0.85)
+
+
+@unittest.skipIf(is_in_ci(), "local test only")
+class TestDeepSeekR10528NVFP4ToMXFP4(TestOnlineQuantizationMemoryLoad):
+    # NVFP4 to MXFP4 online requantization for DeepSeek-R1-0528-NVFP4 on TP=8.
+    # Exercises the MLA attention path (attention_backend=aiter), multi-threaded
+    # weight loading, and the per-expert NVFP4 MoE requantization path.
+    model = "nvidia/DeepSeek-R1-0528-NVFP4"  # NVFP4 model
+    tp = 8
+    runner_args = [
+        "--attention-backend",
+        "aiter",
+        "--model-loader-extra-config",
+        '{"enable_multithread_load": true}',
+    ]
+
+    def test_gsm8k(self):
+        # Requantized NVFP4 -> MXFP4 observed accuracy: ~0.95.
+        self._test_gsm8k(accuracy_threshold=0.90)
+
+
 class TestFP8ToMXFP4DenseTP1(TestOnlineQuantizationMemoryLoad):
     tp = 1
     model = "Qwen/Qwen3-8B-FP8"