From 62af9f67aefec0ed0a0897ecb3c8e58f46e5581c Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Sun, 23 Aug 2026 16:38:33 -0700
Subject: [PATCH 1/7] feat(refit): expose mutable MXFP8 destinations

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 .../inference/quantization/mxfp8_tensor.py    | 31 ++++++++++++++++++-
 megatron/core/inference/quantization/utils.py |  4 +--
 megatron/core/resharding/__init__.py          |  3 +-
 megatron/core/resharding/execution.py         | 14 ++++++---
 megatron/core/resharding/transforms.py        |  4 +--
 .../unit_tests/inference/test_mxfp8_utils.py  | 25 +++++++++++++++
 .../resharding/test_cache_refresh.py          | 11 ++++++-
 7 files changed, 78 insertions(+), 14 deletions(-)

diff --git a/megatron/core/inference/quantization/mxfp8_tensor.py b/megatron/core/inference/quantization/mxfp8_tensor.py
index 702fa4a8ed7..dbfebdf356b 100644
--- a/megatron/core/inference/quantization/mxfp8_tensor.py
+++ b/megatron/core/inference/quantization/mxfp8_tensor.py
@@ -122,6 +122,17 @@ class MXFP8Tensor:
     data: torch.Tensor  # [M, K] fp8_e4m3fn
     scale: torch.Tensor  # 1D swizzled or [M, K // 32] unswizzled scales
     backend: Optional[MXFP8Backend] = None  # quantization and GEMM backend
+    dtype: torch.dtype = torch.bfloat16  # logical, unquantized dtype
+
+    @property
+    def shape(self) -> torch.Size:
+        """Logical tensor shape."""
+        return self.data.shape
+
+    @property
+    def device(self) -> torch.device:
+        """Device holding the quantized tensor."""
+        return self.data.device
 
     def size(self, idx: Optional[int] = None):
         """Wrapper for calling self.data.size()"""
@@ -141,6 +152,24 @@ def scale_2d(self, K: Optional[int] = None) -> torch.Tensor:
         padded_cols = n_col_blocks * MXFP8_SCALE_COL_BLOCK
         return self.scale.reshape(-1, padded_cols)
 
+    def quantize_(self, value: torch.Tensor) -> "MXFP8Tensor":
+        """Quantize a logical tensor into the existing storage."""
+        if value.shape != self.shape:
+            raise ValueError(
+                f"MXFP8 shape mismatch: expected {tuple(self.shape)}, got {tuple(value.shape)}."
+            )
+        if self.backend is None:
+            raise ValueError("Cannot update an MXFP8Tensor without a quantization backend.")
+        value = value.to(device=self.device, dtype=self.dtype).contiguous()
+        quantized = MXFP8Tensor.from_bf16(value, backend=self.backend)
+        self.data.copy_(quantized.data)
+        self.scale.view(torch.uint8).copy_(quantized.scale.view(torch.uint8))
+        return self
+
+    def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
+        """Copy a logical tensor while retaining CUDA-graph-visible storage."""
+        return self.quantize_(value)
+
     @classmethod
     def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend = "flashinfer"):
         """Quantize BF16 tensor to MXFP8.
@@ -168,4 +197,4 @@ def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend
                 f"Unknown MXFP8 quantization backend: '{backend}'. "
                 "Must be 'triton' or 'flashinfer'."
             )
-        return cls(data=data, scale=scale, backend=backend)
+        return cls(data=data, scale=scale, backend=backend, dtype=x.dtype)
diff --git a/megatron/core/inference/quantization/utils.py b/megatron/core/inference/quantization/utils.py
index 504d999e858..9948d00ddfa 100644
--- a/megatron/core/inference/quantization/utils.py
+++ b/megatron/core/inference/quantization/utils.py
@@ -241,9 +241,7 @@ def quantize_params_to_mxfp8(
                     expected_backend=backend,
                     tensor_name=f"persistent MXFP8 parameter {fqn!r}",
                 )
-                new_tensor = MXFP8Tensor.from_bf16(bf16_data, backend=backend)
-                persistent_tensor.data.copy_(new_tensor.data)
-                persistent_tensor.scale.view(torch.uint8).copy_(new_tensor.scale.view(torch.uint8))
+                persistent_tensor.copy_(bf16_data)
                 mcore_tensor = persistent_tensor
             else:
                 # First call: create new MXFP8Tensor
diff --git a/megatron/core/resharding/__init__.py b/megatron/core/resharding/__init__.py
index 7fc122c2400..664d925be0e 100644
--- a/megatron/core/resharding/__init__.py
+++ b/megatron/core/resharding/__init__.py
@@ -1,5 +1,5 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-from .execution import execute_reshard_plan
+from .execution import execute_reshard_plan, refresh_module_caches
 from .planner import (
     build_centralized_reshard_plan,
     build_local_reshard_plan,
@@ -21,6 +21,7 @@
     "build_plan_from_rosters",
     "index_metadata_rosters",
     "execute_reshard_plan",
+    "refresh_module_caches",
     "MXFP8ReshardTransform",
     "ReshardTransform",
     "swap_model_weights",
diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py
index fd28d8cc39f..3965bdde54d 100644
--- a/megatron/core/resharding/execution.py
+++ b/megatron/core/resharding/execution.py
@@ -18,13 +18,17 @@
 logger = logging.getLogger(__name__)
 
 
-def _refresh_module_caches(dst_module: Optional[torch.nn.Module]) -> None:
+def refresh_module_caches(
+    dst_module: torch.nn.Module | list[torch.nn.Module] | tuple[torch.nn.Module, ...] | None,
+) -> None:
     """Refresh parameter-derived caches in the destination module."""
     if dst_module is None:
         return
-    for module in dst_module.modules():
-        if isinstance(module, MegatronModule):
-            module.refresh_cache()
+    roots = dst_module if isinstance(dst_module, (list, tuple)) else (dst_module,)
+    for root in roots:
+        for module in root.modules():
+            if isinstance(module, MegatronModule):
+                module.refresh_cache()
 
 
 @dataclass
@@ -253,7 +257,7 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
             dst_param.quantize_(full_bf16)
     pending_quantized.clear()
 
-    _refresh_module_caches(dst_module)
+    refresh_module_caches(dst_module)
 
     # Ensure all writeback and cache-refresh copies are visible to subsequent CUDA
     # ops (e.g. CUDA graph warmup).  The synchronize() above fires *before* the
diff --git a/megatron/core/resharding/transforms.py b/megatron/core/resharding/transforms.py
index ed5c766d9e8..8f43d3a63cc 100644
--- a/megatron/core/resharding/transforms.py
+++ b/megatron/core/resharding/transforms.py
@@ -280,9 +280,7 @@ def finalize_recv(self, param_name, dst_slice, recv_buffers):
                         f"1D-scale param {param_name!r}: received {written} elements, "
                         f"expected {buf.data.numel()} (duplicate or missing slices?)"
                     )
-                mxfp8 = MXFP8Tensor.from_bf16(accum, backend=buf.backend)
-                buf.data.copy_(mxfp8.data)
-                buf.scale.view(torch.uint8).copy_(mxfp8.scale.view(torch.uint8))
+                buf.copy_(accum)
                 del self._pending_1d[buf_key]
             else:
                 self._pending_1d[buf_key][1] = written
diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py
index 595651f3ad1..8a598473aec 100644
--- a/tests/unit_tests/inference/test_mxfp8_utils.py
+++ b/tests/unit_tests/inference/test_mxfp8_utils.py
@@ -269,6 +269,31 @@ def test_restore_uint8_scale_dtype_after_resharding(self):
         assert scale.dtype == torch.float8_e8m0fnu
         torch.testing.assert_close(scale.view(torch.uint8), scale_bytes)
 
+    @pytest.mark.parametrize("backend", ["triton", "flashinfer"])
+    def test_copy_preserves_storage_and_logical_metadata(self, backend):
+        from megatron.core.inference.quantization.mxfp8_tensor import HAVE_FLASHINFER, MXFP8Tensor
+
+        if backend == "flashinfer" and not HAVE_FLASHINFER:
+            pytest.skip("FlashInfer not available")
+
+        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+        tensor = MXFP8Tensor.from_bf16(source, backend=backend)
+        data_ptr = tensor.data.data_ptr()
+        scale_ptr = tensor.scale.data_ptr()
+
+        assert tensor.shape == source.shape
+        assert tensor.dtype == source.dtype
+        assert tensor.device == source.device
+
+        updated = torch.randn_like(source)
+        assert tensor.copy_(updated) is tensor
+        expected = MXFP8Tensor.from_bf16(updated, backend=backend)
+
+        assert tensor.data.data_ptr() == data_ptr
+        assert tensor.scale.data_ptr() == scale_ptr
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+
     def test_reject_invalid_scale_dtype(self):
         from megatron.core.inference.quantization.mxfp8_tensor import ensure_mxfp8_scale_dtype
 
diff --git a/tests/unit_tests/resharding/test_cache_refresh.py b/tests/unit_tests/resharding/test_cache_refresh.py
index 4d6923abb71..f1231eca937 100644
--- a/tests/unit_tests/resharding/test_cache_refresh.py
+++ b/tests/unit_tests/resharding/test_cache_refresh.py
@@ -4,7 +4,7 @@
 
 import torch
 
-from megatron.core.resharding.execution import execute_reshard_plan
+from megatron.core.resharding.execution import execute_reshard_plan, refresh_module_caches
 from megatron.core.resharding.utils import ReshardPlan
 from megatron.core.transformer.module import MegatronModule
 
@@ -24,6 +24,15 @@ def test_default_refresh_cache_is_noop():
     assert module.refresh_cache() is None
 
 
+def test_refresh_module_caches_accepts_model_lists():
+    events: list[str] = []
+
+    refresh_module_caches(None)
+    refresh_module_caches([torch.nn.Sequential(_CacheAwareModule(events)), torch.nn.Sequential()])
+
+    assert events == ["refresh"]
+
+
 def test_execute_reshard_plan_refreshes_caches_before_final_sync(monkeypatch):
     events: list[str] = []
     target_core = torch.nn.Sequential(_CacheAwareModule(events))

From bc6825277e391d35c7b9cc62bcec8e447fb78e7b Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 07:13:22 -0700
Subject: [PATCH 2/7] fix(refit): preserve MXFP8 constructor compatibility

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 megatron/core/inference/moe/activations.py    |  4 +-
 megatron/core/inference/moe/permute.py        |  5 +-
 .../inference/quantization/mxfp8_tensor.py    | 29 ++++++--
 megatron/core/resharding/execution.py         |  6 +-
 megatron/core/transformer/moe/experts.py      | 14 +++-
 .../unit_tests/inference/test_mxfp8_utils.py  | 70 ++++++++++++++++++-
 .../unit_tests/resharding/test_mxfp8_refit.py |  1 +
 7 files changed, 115 insertions(+), 14 deletions(-)

diff --git a/megatron/core/inference/moe/activations.py b/megatron/core/inference/moe/activations.py
index 929fbb7b9e4..4330388ddb5 100644
--- a/megatron/core/inference/moe/activations.py
+++ b/megatron/core/inference/moe/activations.py
@@ -318,4 +318,6 @@ def squared_relu_and_quantize_mxfp8(
         NUM_BLOCKS=NUM_BLOCKS,
     )
 
-    return MXFP8Tensor(data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton")
+    return MXFP8Tensor(
+        data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), dtype=x.dtype, backend="triton"
+    )
diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py
index f7785b49edf..15c228f1171 100644
--- a/megatron/core/inference/moe/permute.py
+++ b/megatron/core/inference/moe/permute.py
@@ -728,6 +728,9 @@ def permute_and_quantize_mxfp8(
     )
 
     permuted_mxfp8 = MXFP8Tensor(
-        data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton"
+        data=out_fp8,
+        scale=out_scale.view(torch.float8_e8m0fnu),
+        dtype=hidden_states.dtype,
+        backend="triton",
     )
     return permuted_mxfp8, permuted_probs, permutation_map, inclusive_expert_offsets
diff --git a/megatron/core/inference/quantization/mxfp8_tensor.py b/megatron/core/inference/quantization/mxfp8_tensor.py
index dbfebdf356b..8bd21d1f72a 100644
--- a/megatron/core/inference/quantization/mxfp8_tensor.py
+++ b/megatron/core/inference/quantization/mxfp8_tensor.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Literal, Optional
 
 import torch
@@ -122,7 +122,9 @@ class MXFP8Tensor:
     data: torch.Tensor  # [M, K] fp8_e4m3fn
     scale: torch.Tensor  # 1D swizzled or [M, K // 32] unswizzled scales
     backend: Optional[MXFP8Backend] = None  # quantization and GEMM backend
-    dtype: torch.dtype = torch.bfloat16  # logical, unquantized dtype
+    # Keyword-only preserves the pre-existing third positional ``backend`` argument.
+    # Legacy direct constructors have unknown logical dtype until their next update.
+    dtype: Optional[torch.dtype] = field(default=None, kw_only=True)
 
     @property
     def shape(self) -> torch.Size:
@@ -153,29 +155,42 @@ def scale_2d(self, K: Optional[int] = None) -> torch.Tensor:
         return self.scale.reshape(-1, padded_cols)
 
     def quantize_(self, value: torch.Tensor) -> "MXFP8Tensor":
-        """Quantize a logical tensor into the existing storage."""
+        """Quantize a logical tensor into existing storage without changing pointers.
+
+        When this destination's logical dtype is known, the source is cast to it,
+        matching ``torch.Tensor.copy_`` dtype-conversion semantics. Legacy instances
+        with unknown dtype preserve the source dtype. Shape broadcasting is
+        unsupported because MXFP8 scale storage has fixed geometry.
+        """
         if value.shape != self.shape:
             raise ValueError(
                 f"MXFP8 shape mismatch: expected {tuple(self.shape)}, got {tuple(value.shape)}."
             )
         if self.backend is None:
             raise ValueError("Cannot update an MXFP8Tensor without a quantization backend.")
-        value = value.to(device=self.device, dtype=self.dtype).contiguous()
+        logical_dtype = self.dtype if self.dtype is not None else value.dtype
+        value = value.to(device=self.device, dtype=logical_dtype).contiguous()
         quantized = MXFP8Tensor.from_bf16(value, backend=self.backend)
         self.data.copy_(quantized.data)
         self.scale.view(torch.uint8).copy_(quantized.scale.view(torch.uint8))
+        self.dtype = logical_dtype
         return self
 
     def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
-        """Copy a logical tensor while retaining CUDA-graph-visible storage."""
+        """Tensor-compatible alias for :meth:`quantize_` used by generic writeback.
+
+        Unlike ``torch.Tensor.copy_``, this method requires an exact shape and
+        does not accept ``non_blocking``; quantization and storage updates are
+        ordered on the current CUDA stream.
+        """
         return self.quantize_(value)
 
     @classmethod
     def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend = "flashinfer"):
-        """Quantize BF16 tensor to MXFP8.
+        """Quantize a floating-point CUDA tensor to MXFP8.
 
         Args:
-            x: [M, K] BF16 tensor on CUDA.
+            x: [M, K] floating-point tensor on CUDA.
             group_size: MXFP8 group size (default 32).
             backend: 'triton' (fused quantize + swizzle Triton kernel) or
                      'flashinfer' (single fused FlashInfer CUDA kernel).
diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py
index 3965bdde54d..9f44d677519 100644
--- a/megatron/core/resharding/execution.py
+++ b/megatron/core/resharding/execution.py
@@ -21,7 +21,11 @@
 def refresh_module_caches(
     dst_module: torch.nn.Module | list[torch.nn.Module] | tuple[torch.nn.Module, ...] | None,
 ) -> None:
-    """Refresh parameter-derived caches in the destination module."""
+    """Refresh parameter-derived caches in the destination module(s).
+
+    Lists and tuples let external refit callers pass virtual-pipeline model
+    chunks directly. ``None`` is a no-op for send-only ranks.
+    """
     if dst_module is None:
         return
     roots = dst_module if isinstance(dst_module, (list, tuple)) else (dst_module,)
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 6b2ccae8466..51344687638 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -1214,6 +1214,7 @@ def _build_concatenated_mxfp8_weights(self):
         for linear_name, buf_name in [('linear_fc1', '_fc1_weight'), ('linear_fc2', '_fc2_weight')]:
             linear = getattr(self, linear_name)
             q_list, s_list = [], []
+            logical_dtypes = set()
             for i in range(self.num_local_experts):
                 w = getattr(linear, f'weight{i}')
                 if isinstance(w, MXFP8Tensor):
@@ -1228,14 +1229,25 @@ def _build_concatenated_mxfp8_weights(self):
                 validate_mxfp8_tensor(
                     mxfp8, expected_backend=backend, tensor_name=f"{linear_name}.weight{i}"
                 )
+                logical_dtypes.add(mxfp8.dtype)
                 q_list.append(mxfp8.data)
                 s_list.append(mxfp8.scale)
 
             stacked_data = torch.stack(q_list, dim=0).contiguous()
             stacked_scale = torch.stack(s_list, dim=0).contiguous()
+            if len(logical_dtypes) != 1:
+                raise RuntimeError(
+                    f"Expected one logical dtype for {linear_name} expert weights, "
+                    f"got {logical_dtypes}."
+                )
+            logical_dtype = next(iter(logical_dtypes))
 
             setattr(
-                self, buf_name, MXFP8Tensor(data=stacked_data, scale=stacked_scale, backend=backend)
+                self,
+                buf_name,
+                MXFP8Tensor(
+                    data=stacked_data, scale=stacked_scale, dtype=logical_dtype, backend=backend
+                ),
             )
 
             # Redirect per-expert weight .data to views into the stacked buffer,
diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py
index 8a598473aec..4a67ff7f78f 100644
--- a/tests/unit_tests/inference/test_mxfp8_utils.py
+++ b/tests/unit_tests/inference/test_mxfp8_utils.py
@@ -294,6 +294,69 @@ def test_copy_preserves_storage_and_logical_metadata(self, backend):
         assert torch.equal(tensor.data, expected.data)
         assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
 
+    def test_copy_rejects_shape_mismatch(self):
+        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
+
+        tensor = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        with pytest.raises(ValueError, match="shape mismatch"):
+            tensor.copy_(torch.randn(32, 128, device="cuda", dtype=torch.bfloat16))
+
+    def test_copy_requires_backend(self):
+        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
+
+        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+        quantized = MXFP8Tensor.from_bf16(source, backend="triton")
+        backendless = MXFP8Tensor(data=quantized.data, scale=quantized.scale, dtype=source.dtype)
+        with pytest.raises(ValueError, match="without a quantization backend"):
+            backendless.copy_(source)
+
+    def test_copy_preserves_source_dtype_for_legacy_constructor(self):
+        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
+
+        initial = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        tensor = MXFP8Tensor(initial.data.clone(), initial.scale.clone(), "triton")
+        source = torch.randn(16, 128, device="cuda", dtype=torch.float16)
+        expected = MXFP8Tensor.from_bf16(source, backend="triton")
+
+        assert tensor.dtype is None
+        tensor.copy_(source)
+        assert tensor.dtype == source.dtype
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale, expected.scale)
+
+    def test_failed_legacy_copy_keeps_dtype_unknown(self):
+        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
+
+        initial = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        tensor = MXFP8Tensor(
+            initial.data.clone(), initial.scale.clone(), "unsupported"  # type: ignore[arg-type]
+        )
+
+        with pytest.raises(ValueError, match="Unknown MXFP8 quantization backend"):
+            tensor.copy_(torch.randn(16, 128, device="cuda", dtype=torch.float16))
+        assert tensor.dtype is None
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+    def test_copy_uses_explicit_logical_dtype(self, dtype):
+        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
+
+        source = torch.randn(16, 128, device="cuda", dtype=dtype)
+        tensor = MXFP8Tensor.from_bf16(source, backend="triton")
+        update = torch.randn(16, 128, device="cuda", dtype=torch.float32)
+        expected = MXFP8Tensor.from_bf16(update.to(dtype), backend="triton")
+
+        assert tensor.dtype == dtype
+        tensor.copy_(update)
+        assert tensor.dtype == dtype
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale, expected.scale)
+
     def test_reject_invalid_scale_dtype(self):
         from megatron.core.inference.quantization.mxfp8_tensor import ensure_mxfp8_scale_dtype
 
@@ -309,9 +372,9 @@ def test_validate_unswizzled_2d_scale_geometry(self):
         data = torch.empty((64, 128), dtype=torch.float8_e4m3fn, device="cuda")
         scale = torch.empty((64, 4), dtype=torch.uint8, device="cuda")
 
-        validate_mxfp8_tensor(MXFP8Tensor(data, scale, backend="flashinfer"))
+        validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16, backend="flashinfer"))
 
-        invalid = MXFP8Tensor(data, scale[:63], backend="flashinfer")
+        invalid = MXFP8Tensor(data, scale[:63], dtype=torch.bfloat16, backend="flashinfer")
         with pytest.raises(ValueError, match="2D scale has shape"):
             validate_mxfp8_tensor(invalid)
 
@@ -325,7 +388,7 @@ def test_missing_backend_error_is_actionable(self):
         scale = torch.empty(512, dtype=torch.uint8, device="cuda")
 
         with pytest.raises(ValueError, match="backend= explicitly"):
-            validate_mxfp8_tensor(MXFP8Tensor(data, scale))
+            validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16))
 
     @pytest.mark.parametrize("M,K", [(16, 128), (64, 256), (128, 2688)])
     def test_from_bf16_triton(self, M, K):
@@ -750,6 +813,7 @@ def stack_weights() -> MXFP8Tensor:
             return MXFP8Tensor(
                 data=torch.stack([weight.data for weight in per_expert]),
                 scale=torch.stack([weight.scale for weight in per_expert]),
+                dtype=torch.bfloat16,
                 backend="triton",
             )
 
diff --git a/tests/unit_tests/resharding/test_mxfp8_refit.py b/tests/unit_tests/resharding/test_mxfp8_refit.py
index 3506dbc1fdd..8845cc65bef 100644
--- a/tests/unit_tests/resharding/test_mxfp8_refit.py
+++ b/tests/unit_tests/resharding/test_mxfp8_refit.py
@@ -188,6 +188,7 @@ def test_mixed_persistent_buffer_backends_are_rejected(self):
         flashinfer_buffer = MXFP8Tensor(
             data=triton_buffer.data.clone(),
             scale=triton_buffer.scale.view(torch.uint8).clone(),
+            dtype=torch.bfloat16,
             backend="flashinfer",
         )
 

From af5b8b0508bf680cdb34764997c4f7d189f8deb3 Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 10:36:02 -0700
Subject: [PATCH 3/7] test(mxfp8): hoist refit test imports

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 tests/unit_tests/inference/test_mxfp8_utils.py | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py
index 4a67ff7f78f..49ca0c15f9a 100644
--- a/tests/unit_tests/inference/test_mxfp8_utils.py
+++ b/tests/unit_tests/inference/test_mxfp8_utils.py
@@ -11,6 +11,8 @@
 import pytest
 import torch
 
+from megatron.core.inference.quantization.mxfp8_tensor import HAVE_FLASHINFER, MXFP8Tensor
+
 pytestmark = [
     pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
     pytest.mark.internal,
@@ -271,8 +273,6 @@ def test_restore_uint8_scale_dtype_after_resharding(self):
 
     @pytest.mark.parametrize("backend", ["triton", "flashinfer"])
     def test_copy_preserves_storage_and_logical_metadata(self, backend):
-        from megatron.core.inference.quantization.mxfp8_tensor import HAVE_FLASHINFER, MXFP8Tensor
-
         if backend == "flashinfer" and not HAVE_FLASHINFER:
             pytest.skip("FlashInfer not available")
 
@@ -295,8 +295,6 @@ def test_copy_preserves_storage_and_logical_metadata(self, backend):
         assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
 
     def test_copy_rejects_shape_mismatch(self):
-        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
-
         tensor = MXFP8Tensor.from_bf16(
             torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
         )
@@ -304,8 +302,6 @@ def test_copy_rejects_shape_mismatch(self):
             tensor.copy_(torch.randn(32, 128, device="cuda", dtype=torch.bfloat16))
 
     def test_copy_requires_backend(self):
-        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
-
         source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
         quantized = MXFP8Tensor.from_bf16(source, backend="triton")
         backendless = MXFP8Tensor(data=quantized.data, scale=quantized.scale, dtype=source.dtype)
@@ -313,8 +309,6 @@ def test_copy_requires_backend(self):
             backendless.copy_(source)
 
     def test_copy_preserves_source_dtype_for_legacy_constructor(self):
-        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
-
         initial = MXFP8Tensor.from_bf16(
             torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
         )
@@ -329,8 +323,6 @@ def test_copy_preserves_source_dtype_for_legacy_constructor(self):
         assert torch.equal(tensor.scale, expected.scale)
 
     def test_failed_legacy_copy_keeps_dtype_unknown(self):
-        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
-
         initial = MXFP8Tensor.from_bf16(
             torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
         )
@@ -344,8 +336,6 @@ def test_failed_legacy_copy_keeps_dtype_unknown(self):
 
     @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
     def test_copy_uses_explicit_logical_dtype(self, dtype):
-        from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
-
         source = torch.randn(16, 128, device="cuda", dtype=dtype)
         tensor = MXFP8Tensor.from_bf16(source, backend="triton")
         update = torch.randn(16, 128, device="cuda", dtype=torch.float32)

From 3b60bfb51bc1fb6d7a020b8a3cec7978166b9aab Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 10:58:34 -0700
Subject: [PATCH 4/7] fix(refit): preserve MXFP8 update precision

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 .../inference/quantization/mxfp8_tensor.py    | 40 ++++++++++++-----
 megatron/core/resharding/execution.py         |  6 ++-
 megatron/core/transformer/moe/experts.py      | 15 ++++---
 .../unit_tests/inference/test_mxfp8_utils.py  | 44 ++++++++++++++++---
 .../unit_tests/resharding/test_mxfp8_refit.py |  4 ++
 5 files changed, 83 insertions(+), 26 deletions(-)

diff --git a/megatron/core/inference/quantization/mxfp8_tensor.py b/megatron/core/inference/quantization/mxfp8_tensor.py
index 8bd21d1f72a..8276cbfcb0d 100644
--- a/megatron/core/inference/quantization/mxfp8_tensor.py
+++ b/megatron/core/inference/quantization/mxfp8_tensor.py
@@ -123,12 +123,15 @@ class MXFP8Tensor:
     scale: torch.Tensor  # 1D swizzled or [M, K // 32] unswizzled scales
     backend: Optional[MXFP8Backend] = None  # quantization and GEMM backend
     # Keyword-only preserves the pre-existing third positional ``backend`` argument.
-    # Legacy direct constructors have unknown logical dtype until their next update.
+    # This metadata supports tensor-like conversion consumers such as Megatron
+    # Bridge; it does not constrain the precision accepted by ``quantize_``/``copy_``.
+    # Legacy direct constructors have unknown logical dtype until their next
+    # successful update.
     dtype: Optional[torch.dtype] = field(default=None, kw_only=True)
 
     @property
     def shape(self) -> torch.Size:
-        """Logical tensor shape."""
+        """Shape of the quantized data storage."""
         return self.data.shape
 
     @property
@@ -157,23 +160,34 @@ def scale_2d(self, K: Optional[int] = None) -> torch.Tensor:
     def quantize_(self, value: torch.Tensor) -> "MXFP8Tensor":
         """Quantize a logical tensor into existing storage without changing pointers.
 
-        When this destination's logical dtype is known, the source is cast to it,
-        matching ``torch.Tensor.copy_`` dtype-conversion semantics. Legacy instances
-        with unknown dtype preserve the source dtype. Shape broadcasting is
-        unsupported because MXFP8 scale storage has fixed geometry.
+        The source dtype is preserved when supported by the configured backend so
+        the quantizer does not introduce an avoidable intermediate downcast.
+        FlashInfer FP32 inputs are converted to BF16 because that backend accepts
+        FP16/BF16 inputs only. For legacy instances with unknown logical dtype, the
+        first successful update records the dtype actually passed to the quantizer.
+        Shape broadcasting is unsupported because MXFP8 scale storage has fixed
+        geometry.
         """
+        if self.data.ndim != 2:
+            raise ValueError(
+                "In-place MXFP8 updates require 2D destination storage; "
+                f"got shape {tuple(self.data.shape)}."
+            )
         if value.shape != self.shape:
             raise ValueError(
                 f"MXFP8 shape mismatch: expected {tuple(self.shape)}, got {tuple(value.shape)}."
             )
         if self.backend is None:
             raise ValueError("Cannot update an MXFP8Tensor without a quantization backend.")
-        logical_dtype = self.dtype if self.dtype is not None else value.dtype
-        value = value.to(device=self.device, dtype=logical_dtype).contiguous()
+        value = value.to(device=self.device)
+        if self.backend == "flashinfer" and value.dtype == torch.float32:
+            value = value.to(dtype=torch.bfloat16)
+        value = value.contiguous()
         quantized = MXFP8Tensor.from_bf16(value, backend=self.backend)
         self.data.copy_(quantized.data)
         self.scale.view(torch.uint8).copy_(quantized.scale.view(torch.uint8))
-        self.dtype = logical_dtype
+        if self.dtype is None:
+            self.dtype = value.dtype
         return self
 
     def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
@@ -181,7 +195,9 @@ def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
 
         Unlike ``torch.Tensor.copy_``, this method requires an exact shape and
         does not accept ``non_blocking``; quantization and storage updates are
-        ordered on the current CUDA stream.
+        ordered on the current CUDA stream. The alias lets external conversion
+        integrations such as Megatron Bridge treat plain and MXFP8 destinations
+        uniformly while ``quantize_`` remains available to existing callers.
         """
         return self.quantize_(value)
 
@@ -189,6 +205,10 @@ def copy_(self, value: torch.Tensor) -> "MXFP8Tensor":
     def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend = "flashinfer"):
         """Quantize a floating-point CUDA tensor to MXFP8.
 
+        The historical method name is retained for compatibility. The Triton
+        backend accepts BF16, FP16, and FP32 inputs; FlashInfer accepts BF16 and
+        FP16 inputs.
+
         Args:
             x: [M, K] floating-point tensor on CUDA.
             group_size: MXFP8 group size (default 32).
diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py
index 9f44d677519..2f8d7bdb45c 100644
--- a/megatron/core/resharding/execution.py
+++ b/megatron/core/resharding/execution.py
@@ -23,8 +23,10 @@ def refresh_module_caches(
 ) -> None:
     """Refresh parameter-derived caches in the destination module(s).
 
-    Lists and tuples let external refit callers pass virtual-pipeline model
-    chunks directly. ``None`` is a no-op for send-only ranks.
+    Lists and tuples let external refit callers, including Megatron Bridge,
+    pass virtual-pipeline model chunks directly. ``None`` is a no-op for
+    send-only ranks. The ``dst_module`` parameter name remains stable for
+    callers that use keyword arguments.
     """
     if dst_module is None:
         return
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 51344687638..2d67726594e 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -1214,7 +1214,7 @@ def _build_concatenated_mxfp8_weights(self):
         for linear_name, buf_name in [('linear_fc1', '_fc1_weight'), ('linear_fc2', '_fc2_weight')]:
             linear = getattr(self, linear_name)
             q_list, s_list = [], []
-            logical_dtypes = set()
+            logical_dtypes: set[torch.dtype | None] = set()
             for i in range(self.num_local_experts):
                 w = getattr(linear, f'weight{i}')
                 if isinstance(w, MXFP8Tensor):
@@ -1233,14 +1233,15 @@ def _build_concatenated_mxfp8_weights(self):
                 q_list.append(mxfp8.data)
                 s_list.append(mxfp8.scale)
 
-            stacked_data = torch.stack(q_list, dim=0).contiguous()
-            stacked_scale = torch.stack(s_list, dim=0).contiguous()
-            if len(logical_dtypes) != 1:
+            known_dtypes = logical_dtypes - {None}
+            if len(known_dtypes) > 1:
                 raise RuntimeError(
-                    f"Expected one logical dtype for {linear_name} expert weights, "
-                    f"got {logical_dtypes}."
+                    f"Conflicting logical dtypes for {linear_name} expert weights: "
+                    f"{known_dtypes}."
                 )
-            logical_dtype = next(iter(logical_dtypes))
+            logical_dtype = next(iter(known_dtypes), None)
+            stacked_data = torch.stack(q_list, dim=0).contiguous()
+            stacked_scale = torch.stack(s_list, dim=0).contiguous()
 
             setattr(
                 self,
diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py
index 49ca0c15f9a..1a36bd9f007 100644
--- a/tests/unit_tests/inference/test_mxfp8_utils.py
+++ b/tests/unit_tests/inference/test_mxfp8_utils.py
@@ -334,18 +334,48 @@ def test_failed_legacy_copy_keeps_dtype_unknown(self):
             tensor.copy_(torch.randn(16, 128, device="cuda", dtype=torch.float16))
         assert tensor.dtype is None
 
-    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
-    def test_copy_uses_explicit_logical_dtype(self, dtype):
-        source = torch.randn(16, 128, device="cuda", dtype=dtype)
+    def test_copy_does_not_downcast_update_to_logical_dtype(self):
+        source = torch.randn(16, 128, device="cuda", dtype=torch.float16)
         tensor = MXFP8Tensor.from_bf16(source, backend="triton")
+        update = torch.linspace(
+            -100_000, 100_000, steps=16 * 128, device="cuda", dtype=torch.float32
+        ).reshape(16, 128)
+        expected = MXFP8Tensor.from_bf16(update, backend="triton")
+
+        assert tensor.dtype == torch.float16
+        tensor.copy_(update)
+        assert tensor.dtype == torch.float16
+        assert torch.equal(tensor.data, expected.data)
+        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+
+    def test_copy_normalizes_flashinfer_fp32_input_to_bf16(self):
+        if not HAVE_FLASHINFER:
+            pytest.skip("FlashInfer not available")
+
+        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+        tensor = MXFP8Tensor.from_bf16(source, backend="flashinfer")
         update = torch.randn(16, 128, device="cuda", dtype=torch.float32)
-        expected = MXFP8Tensor.from_bf16(update.to(dtype), backend="triton")
+        expected = MXFP8Tensor.from_bf16(update.to(torch.bfloat16), backend="flashinfer")
 
-        assert tensor.dtype == dtype
         tensor.copy_(update)
-        assert tensor.dtype == dtype
+
+        assert tensor.dtype == torch.bfloat16
         assert torch.equal(tensor.data, expected.data)
-        assert torch.equal(tensor.scale, expected.scale)
+        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+
+    def test_copy_rejects_stacked_storage(self):
+        tensor = MXFP8Tensor.from_bf16(
+            torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton"
+        )
+        stacked = MXFP8Tensor(
+            data=torch.stack([tensor.data, tensor.data]),
+            scale=torch.stack([tensor.scale, tensor.scale]),
+            dtype=tensor.dtype,
+            backend=tensor.backend,
+        )
+
+        with pytest.raises(ValueError, match="require 2D destination storage"):
+            stacked.copy_(torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16))
 
     def test_reject_invalid_scale_dtype(self):
         from megatron.core.inference.quantization.mxfp8_tensor import ensure_mxfp8_scale_dtype
diff --git a/tests/unit_tests/resharding/test_mxfp8_refit.py b/tests/unit_tests/resharding/test_mxfp8_refit.py
index 8845cc65bef..a97511ff08c 100644
--- a/tests/unit_tests/resharding/test_mxfp8_refit.py
+++ b/tests/unit_tests/resharding/test_mxfp8_refit.py
@@ -219,6 +219,8 @@ class Namespace:
                 tensor = MXFP8Tensor.from_bf16(
                     torch.randn(M, K, dtype=torch.bfloat16, device="cuda"), backend="triton"
                 )
+                if expert_idx == 0:
+                    tensor.dtype = None  # Simulate a legacy direct constructor.
                 setattr(linear, f"weight{expert_idx}", tensor)
                 buffers[f"{linear_name}.weight{expert_idx}"] = tensor
 
@@ -229,6 +231,8 @@ class Namespace:
         assert not torch.is_inference(grouped_mlp._fc1_weight.scale)
         assert not torch.is_inference(grouped_mlp._fc2_weight.data)
         assert not torch.is_inference(grouped_mlp._fc2_weight.scale)
+        assert grouped_mlp._fc1_weight.dtype == torch.bfloat16
+        assert grouped_mlp._fc2_weight.dtype == torch.bfloat16
 
         transform = MXFP8ReshardTransform(
             convertible_params=set(buffers), persistent_buffers=buffers

From f64018c9f1a3737353f6ff0b3b12207c881b3e0c Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 11:47:23 -0700
Subject: [PATCH 5/7] refactor(moe): clarify MXFP8 source dtype propagation

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 megatron/core/transformer/moe/experts.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 2d67726594e..9a295b25b02 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -1214,7 +1214,7 @@ def _build_concatenated_mxfp8_weights(self):
         for linear_name, buf_name in [('linear_fc1', '_fc1_weight'), ('linear_fc2', '_fc2_weight')]:
             linear = getattr(self, linear_name)
             q_list, s_list = [], []
-            logical_dtypes: set[torch.dtype | None] = set()
+            source_dtype: torch.dtype | None = None
             for i in range(self.num_local_experts):
                 w = getattr(linear, f'weight{i}')
                 if isinstance(w, MXFP8Tensor):
@@ -1229,17 +1229,16 @@ def _build_concatenated_mxfp8_weights(self):
                 validate_mxfp8_tensor(
                     mxfp8, expected_backend=backend, tensor_name=f"{linear_name}.weight{i}"
                 )
-                logical_dtypes.add(mxfp8.dtype)
+                if mxfp8.dtype is not None:
+                    if source_dtype is not None and mxfp8.dtype != source_dtype:
+                        raise RuntimeError(
+                            f"Conflicting source dtypes for {linear_name} expert weights: "
+                            f"{source_dtype} and {mxfp8.dtype}."
+                        )
+                    source_dtype = mxfp8.dtype
                 q_list.append(mxfp8.data)
                 s_list.append(mxfp8.scale)
 
-            known_dtypes = logical_dtypes - {None}
-            if len(known_dtypes) > 1:
-                raise RuntimeError(
-                    f"Conflicting logical dtypes for {linear_name} expert weights: "
-                    f"{known_dtypes}."
-                )
-            logical_dtype = next(iter(known_dtypes), None)
             stacked_data = torch.stack(q_list, dim=0).contiguous()
             stacked_scale = torch.stack(s_list, dim=0).contiguous()
 
@@ -1247,7 +1246,7 @@ def _build_concatenated_mxfp8_weights(self):
                 self,
                 buf_name,
                 MXFP8Tensor(
-                    data=stacked_data, scale=stacked_scale, dtype=logical_dtype, backend=backend
+                    data=stacked_data, scale=stacked_scale, dtype=source_dtype, backend=backend
                 ),
             )
 

From f8c9186b0be2150cc0bba4d20084fbbbadbd032c Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 13:54:39 -0700
Subject: [PATCH 6/7] fix(refit): refresh caches after native execution

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 megatron/core/resharding/execution.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py
index 2f8d7bdb45c..aec33a94616 100644
--- a/megatron/core/resharding/execution.py
+++ b/megatron/core/resharding/execution.py
@@ -113,7 +113,7 @@ def execute_reshard_plan(
         torch.cuda.synchronize()
         if service.requires_process_group_barrier:
             dist.barrier(group=group)
-        _refresh_module_caches(dst_module)
+        refresh_module_caches(dst_module)
         torch.cuda.synchronize()
         logger.info("Reshard complete")
         return

From 096b22f53353fd5aa762188573d0a6b643b6dc2e Mon Sep 17 00:00:00 2001
From: wdykas <wdykas@nvidia.com>
Date: Tue, 25 Aug 2026 17:30:51 -0700
Subject: [PATCH 7/7] test(mxfp8): stabilize dtype normalization coverage

Signed-off-by: wdykas <wdykas@nvidia.com>
---
 megatron/core/transformer/moe/experts.py      |  4 +--
 .../unit_tests/inference/test_mxfp8_utils.py  | 27 ++++++++++++-------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 9a295b25b02..2b068ff1107 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -1230,12 +1230,12 @@ def _build_concatenated_mxfp8_weights(self):
                     mxfp8, expected_backend=backend, tensor_name=f"{linear_name}.weight{i}"
                 )
                 if mxfp8.dtype is not None:
-                    if source_dtype is not None and mxfp8.dtype != source_dtype:
+                    source_dtype = source_dtype or mxfp8.dtype
+                    if mxfp8.dtype != source_dtype:
                         raise RuntimeError(
                             f"Conflicting source dtypes for {linear_name} expert weights: "
                             f"{source_dtype} and {mxfp8.dtype}."
                         )
-                    source_dtype = mxfp8.dtype
                 q_list.append(mxfp8.data)
                 s_list.append(mxfp8.scale)
 
diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py
index 1a36bd9f007..fa05ecb1ec7 100644
--- a/tests/unit_tests/inference/test_mxfp8_utils.py
+++ b/tests/unit_tests/inference/test_mxfp8_utils.py
@@ -348,20 +348,29 @@ def test_copy_does_not_downcast_update_to_logical_dtype(self):
         assert torch.equal(tensor.data, expected.data)
         assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
 
-    def test_copy_normalizes_flashinfer_fp32_input_to_bf16(self):
-        if not HAVE_FLASHINFER:
-            pytest.skip("FlashInfer not available")
-
-        source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
-        tensor = MXFP8Tensor.from_bf16(source, backend="flashinfer")
+    def test_copy_normalizes_flashinfer_fp32_input_to_bf16(self, monkeypatch):
+        tensor = MXFP8Tensor(
+            data=torch.empty((16, 128), device="cuda", dtype=torch.float8_e4m3fn),
+            scale=torch.empty(512, device="cuda", dtype=torch.uint8),
+            backend="flashinfer",
+        )
         update = torch.randn(16, 128, device="cuda", dtype=torch.float32)
-        expected = MXFP8Tensor.from_bf16(update.to(torch.bfloat16), backend="flashinfer")
+        expected_data = torch.zeros_like(tensor.data)
+        expected_scale = torch.zeros_like(tensor.scale)
+        quantizer_input = {}
+
+        def fake_from_bf16(cls, value, group_size=32, backend="flashinfer"):
+            quantizer_input["dtype"] = value.dtype
+            return cls(data=expected_data, scale=expected_scale, backend=backend, dtype=value.dtype)
+
+        monkeypatch.setattr(MXFP8Tensor, "from_bf16", classmethod(fake_from_bf16))
 
         tensor.copy_(update)
 
+        assert quantizer_input["dtype"] == torch.bfloat16
         assert tensor.dtype == torch.bfloat16
-        assert torch.equal(tensor.data, expected.data)
-        assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8))
+        assert torch.equal(tensor.data, expected_data)
+        assert torch.equal(tensor.scale, expected_scale)
 
     def test_copy_rejects_stacked_storage(self):
         tensor = MXFP8Tensor.from_bf16(