diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py
index ce2eecb678..70479d575f 100644
--- a/src/maxtext/checkpoint_conversion/to_maxtext.py
+++ b/src/maxtext/checkpoint_conversion/to_maxtext.py
@@ -53,6 +53,10 @@
 from functools import partial
 import json
 import os
+
+if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
+  os.environ["HF_HOME"] = "/dev/shm/hf_cache"
+
 import sys
 import threading
 import time
@@ -83,6 +87,10 @@
 except ImportError:
   torch = None
 
+for dtype_name in ("float8_e4m3fn", "float8_e5m2", "bfloat16"):
+  if hasattr(ml_dtypes, dtype_name) and not hasattr(np, dtype_name):
+    setattr(np, dtype_name, getattr(ml_dtypes, dtype_name))
+
 
 absl.logging.set_verbosity(absl.logging.INFO)  # for max_logging.log
 
@@ -178,7 +186,25 @@ def get_tensor(self, key: str) -> np.ndarray:
     and reads only the required tensor's data from disk.
     """
     # Handle single-file models (shard map key might be None or we just know the filename)
-    shard_name = self.shard_map.get(key)
+    resolved_key = key
+    shard_name = self.shard_map.get(resolved_key)
+    if shard_name is None:
+      # Check fallback for .weight_scale vs .scale and inverse scales
+      if resolved_key.endswith(".weight_scale"):
+        for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+          alt_key = resolved_key[:-len(".weight_scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+      elif resolved_key.endswith(".scale"):
+        for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+          alt_key = resolved_key[:-len(".scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+
     if shard_name is None and None in self.shard_map:
       shard_name = self.shard_map[None]
     elif shard_name is None:
@@ -205,8 +231,35 @@ def get_tensor(self, key: str) -> np.ndarray:
     # STEP 2: Lock ONLY the reading into RAM.
     # This prevents multiple threads from simultaneously allocating large chunks of RAM.
     with self._ram_lock:
-      with safe_open(local_path, framework="np", device="cpu") as f:
-        return f.get_tensor(key)
+      framework = "pt" if torch is not None else "np"
+      with safe_open(local_path, framework=framework, device="cpu") as f:
+        final_key = resolved_key
+        if final_key not in f.keys():
+          if final_key.endswith(".weight_scale"):
+            for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+              alt_key = final_key[:-len(".weight_scale")] + suffix
+              if alt_key in f.keys():
+                final_key = alt_key
+                break
+          elif final_key.endswith(".scale"):
+            for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+              alt_key = final_key[:-len(".scale")] + suffix
+              if alt_key in f.keys():
+                final_key = alt_key
+                break
+        t = f.get_tensor(final_key)
+        if torch is not None and isinstance(t, torch.Tensor):
+          if hasattr(torch, "float8_e4m3fn") and t.dtype == torch.float8_e4m3fn:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e4m3fn)
+          elif hasattr(torch, "float8_e5m2") and t.dtype == torch.float8_e5m2:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e5m2)
+          elif t.dtype == torch.bfloat16:
+            return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
+          elif t.dtype == torch.float16:
+            return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
+          else:
+            return t.numpy()
+        return t
 
 
 class LazyTensor:
@@ -224,7 +277,18 @@ def __init__(
   ):
     self._load_fn = load_fn
     self.shape = shape
-    self.dtype = np.dtype(dtype)
+    try:
+      self.dtype = np.dtype(dtype)
+    except (TypeError, ValueError):
+      dtype_str = str(dtype)
+      if "float8_e4m3fn" in dtype_str:
+        self.dtype = np.dtype(ml_dtypes.float8_e4m3fn)
+      elif "float8_e5m2" in dtype_str:
+        self.dtype = np.dtype(ml_dtypes.float8_e5m2)
+      elif "bfloat16" in dtype_str:
+        self.dtype = np.dtype(ml_dtypes.bfloat16)
+      else:
+        self.dtype = np.dtype(np.float32)
     self.ndim = len(shape)
     self.name = name
 
@@ -433,7 +497,7 @@ def _build_single_axis_stacked_tensor(
 
   if config.scan_layers:
     # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0
   else:
     # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0).
     axis_to_stack = 0
@@ -732,7 +796,7 @@ def convert_lora_to_maxtext_adapter(
 
   mt_adapter_tree = {}
   mapped_count = 0
-  target_dtype = ml_dtypes.bfloat16 if save_dtype == "bfloat16" else np.float32
+  target_dtype = ml_dtypes.bfloat16 if save_dtype in ("bfloat16", "float8_e4m3fn", "float8_e5m2") else np.float32
 
   collected_weights = {}
 
@@ -867,6 +931,17 @@ def main(
     simulated_cpu_devices_count: int = 16,
 ) -> None:
   overall_start = time.time()
+  cleaned_args = []
+  for arg in args:
+    if arg.startswith("save_dtype="):
+      save_dtype = arg.split("=", 1)[1]
+    elif arg.startswith("hf_model_path="):
+      hf_model_path = arg.split("=", 1)[1]
+    elif arg.startswith("lazy_load_tensors="):
+      lazy_load_tensors = str2bool(arg.split("=", 1)[1])
+    else:
+      cleaned_args.append(arg)
+  args = cleaned_args
   # Check if the user is using an Instruct version. If so, use the base model architecture
   for i, arg in enumerate(args):
     if arg.startswith("model_name="):
@@ -983,21 +1058,62 @@ def main(
           }
 
       def _eager_getter(key):
-        if key not in hf_state_dict_numpy:
+        resolved_key = key
+        if resolved_key not in hf_state_dict_numpy:
+          if resolved_key.endswith(".weight_scale"):
+            for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+              alt_key = resolved_key[:-len(".weight_scale")] + suffix
+              if alt_key in hf_state_dict_numpy:
+                resolved_key = alt_key
+                break
+          elif resolved_key.endswith(".scale"):
+            for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+              alt_key = resolved_key[:-len(".scale")] + suffix
+              if alt_key in hf_state_dict_numpy:
+                resolved_key = alt_key
+                break
+        if resolved_key not in hf_state_dict_numpy:
           raise ValueError(f"HuggingFace key {key} not found in state_dict.")
-        v = hf_state_dict_numpy[key]
+
+        v = hf_state_dict_numpy[resolved_key]
         # target dtype is "float32"
-        if save_dtype == DType.FLOAT32:
-          return v.to(torch.float32).numpy()
+        if save_dtype == DType.FLOAT32 or save_dtype == "float32":
+          if torch is not None and isinstance(v, torch.Tensor):
+            return v.to(torch.float32).cpu().numpy()
+          return np.asarray(v, dtype=np.float32)
         # target dtype is "bfloat16"
-        elif save_dtype == DType.BFLOAT16:
+        elif save_dtype == DType.BFLOAT16 or save_dtype == "bfloat16":
           # - torch.bfloat16 -> torch.float32 -> np.float32 -> ml_dtypes.bfloat16
           #   As numpy doesn't accept bfloat16 directly, we convert to float32 first
           # - torch.float16 -> np.float16 -> ml_dtypes.bfloat16
           # - torch.float32 -> np.float32 -> ml_dtypes.bfloat16
-          if v.dtype == torch.bfloat16:
-            v = v.to(torch.float32)
-          return v.numpy().astype(ml_dtypes.bfloat16)
+          if torch is not None and isinstance(v, torch.Tensor):
+            if v.dtype == torch.bfloat16:
+              v = v.to(torch.float32)
+            return v.to(torch.float32).cpu().numpy().astype(ml_dtypes.bfloat16)
+          return np.asarray(v).astype(ml_dtypes.bfloat16)
+        # target dtype is "float8_e4m3fn" or "float8_e5m2"
+        elif save_dtype in (
+            DType.FLOAT8_E4M3FN,
+            "float8_e4m3fn",
+            DType.FLOAT8_E5M2,
+            "float8_e5m2",
+        ):
+          if torch is not None and isinstance(v, torch.Tensor):
+            if hasattr(torch, "float8_e4m3fn") and v.dtype == torch.float8_e4m3fn:
+              return v.view(torch.uint8).cpu().numpy().view(ml_dtypes.float8_e4m3fn)
+            elif hasattr(torch, "float8_e5m2") and v.dtype == torch.float8_e5m2:
+              return v.view(torch.uint8).cpu().numpy().view(ml_dtypes.float8_e5m2)
+            elif v.dtype == torch.bfloat16:
+              return v.to(torch.float32).cpu().numpy().astype(ml_dtypes.bfloat16)
+            elif v.dtype == torch.float16:
+              return v.to(torch.float32).cpu().numpy().astype(ml_dtypes.bfloat16)
+            elif v.dtype == torch.float32:
+              return v.cpu().numpy()
+            else:
+              return v.to(torch.float32).cpu().numpy()
+          else:
+            return np.asarray(v)
         raise NotImplementedError(f"Save dtype {save_dtype} is not currently implemented.")
 
       tensor_getter = _eager_getter
@@ -1160,7 +1276,7 @@ def _eager_getter(key):
       type=str,
       required=False,
       default="bfloat16",
-      choices=["float32", "bfloat16"],
+      choices=["float32", "bfloat16", "float8_e4m3fn", "float8_e5m2"],
       help="Save MaxText weights in specified dtype",
   )
   # Determines the logical sharding of the output checkpoint by partitioning
@@ -1179,6 +1295,23 @@ def _eager_getter(key):
   parser.add_argument(
       "--simulated_cpu_devices_count", type=int, required=False, default=16, help="Sharding of checkpoint"
   )
+  # Normalize key=value CLI arguments for local_args if passed without leading dashes
+  normalized_argv = [sys.argv[0]]
+  for arg in sys.argv[1:]:
+    for prefix in (
+        "save_dtype=",
+        "hf_model_path=",
+        "lazy_load_tensors=",
+        "eager_load_method=",
+        "revision=",
+        "simulated_cpu_devices_count=",
+    ):
+      if arg.startswith(prefix):
+        arg = "--" + arg
+        break
+    normalized_argv.append(arg)
+  sys.argv = normalized_argv
+
   # Parse local arguments
   # Parse known args returns the namespace AND the list of remaining arguments
   local_args, remaining_args = parser.parse_known_args()
diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py
index 89abd56d4c..97137f00d9 100644
--- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py
+++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py
@@ -1930,7 +1930,10 @@ def __init__(self, **kwargs):
     "qwen3-omni-30b-a3b": qwen3_omni_30b_a3b_config,
     "qwen3-next-80b-a3b": qwen3_next_80b_a3b_config,
     "qwen3.5-397b-a17b": qwen3_5_397b_a17b_config,
+    "qwen3.5-397b-a17b-fp8": qwen3_5_397b_a17b_config,
     "qwen3.5-35b-a3b": qwen3_5_35b_a3b_config,
+    "qwen3.5-35b-a3b-fp8": qwen3_5_35b_a3b_config,
+    "qwen3.5-35b-fp8": qwen3_5_35b_a3b_config,
     "mixtral-8x7b": mixtral_8x7b_config,
     "mixtral-8x22b": mixtral_8x22b_config,
     "olmo3-7b": olmo3_7b_config,
diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py
index 85dd1d6ea0..c80c643d24 100644
--- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py
+++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py
@@ -1317,6 +1317,9 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config):
     "mixtral-8x7b": MIXTRAL_HF_WEIGHTS_TO_SHAPE,
     "mixtral-8x22b": MIXTRAL_HF_WEIGHTS_TO_SHAPE,
     "qwen3.5-35b-a3b": QWEN3_5_HF_WEIGHTS_TO_SHAPE,
+    "qwen3.5-35b-a3b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE,
+    "qwen3.5-35b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE,
     "qwen3.5-397b-a17b": QWEN3_5_HF_WEIGHTS_TO_SHAPE,
+    "qwen3.5-397b-a17b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE,
     "qwen3-next-80b-a3b": QWEN3_NEXT_HF_WEIGHTS_TO_SHAPE,
 }
diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
index 94e96173f1..62772efd08 100644
--- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py
+++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
@@ -872,6 +872,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
   """
   num_main_layers = config["text_config"]["num_hidden_layers"]
   layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
+  is_quantized = getattr(maxtext_config, "weight_dtype", None) in (
+      "float8_e4m3fn",
+      "float8_e5m2",
+  )
 
   # 1. Non-layer specific weight mappings
   mapping = {
@@ -920,6 +924,23 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
                 ],
             }
         )
+        if is_quantized:
+          mapping.update(
+              {
+                  f"{prefix}-attention-attention-query-kernel_scale": [
+                      f"model.language_model.layers.{i}.self_attn.q_proj.weight_scale_inv" for i in hf_indices
+                  ],
+                  f"{prefix}-attention-attention-key-kernel_scale": [
+                      f"model.language_model.layers.{i}.self_attn.k_proj.weight_scale_inv" for i in hf_indices
+                  ],
+                  f"{prefix}-attention-attention-value-kernel_scale": [
+                      f"model.language_model.layers.{i}.self_attn.v_proj.weight_scale_inv" for i in hf_indices
+                  ],
+                  f"{prefix}-attention-attention-out-kernel_scale": [
+                      f"model.language_model.layers.{i}.self_attn.o_proj.weight_scale_inv" for i in hf_indices
+                  ],
+              }
+          )
       else:
         # Linear/Hybrid Attention Block
         mapping.update(  # pyrefly: ignore[no-matching-overload]
@@ -955,6 +976,21 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
                 ],
             }
         )
+        if is_quantized:
+          mapping.update(
+              {
+                  f"{prefix}-attention-in_proj_qkvz-kernel_scale": [
+                      (
+                          f"model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight_scale_inv",
+                          f"model.language_model.layers.{i}.linear_attn.in_proj_z.weight_scale_inv",
+                      )
+                      for i in hf_indices
+                  ],
+                  f"{prefix}-attention-out_proj-kernel_scale": [
+                      f"model.language_model.layers.{i}.linear_attn.out_proj.weight_scale_inv" for i in hf_indices
+                  ],
+              }
+          )
 
       # 3. Handle MLP: Gates and Shared Experts
       mapping.update(  # pyrefly: ignore[no-matching-overload]
@@ -976,18 +1012,63 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
               ],
           }
       )
+      if is_quantized:
+        mapping.update(
+            {
+                f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight_scale_inv" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight_scale_inv" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert-wo-kernel_scale": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight_scale_inv" for i in hf_indices
+                ],
+            }
+        )
 
       # 4. Handle MoE Routed Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-wo": [
-                  f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices
-              ],
-              (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [
-                  f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices
-              ],
-          }
-      )
+      if is_quantized:
+        num_experts = config.get("text_config", config).get("num_experts", 256)
+        mapping.update(
+            {
+                f"{prefix}-mlp-routed_experts-wi_0": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_0_scale": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_1": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_1_scale": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wo": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wo_scale": [
+                    [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" for i in hf_indices]
+                    for e in range(num_experts)
+                ],
+            }
+        )
+      else:
+        mapping.update(  # pyrefly: ignore[no-matching-overload]
+            {
+                f"{prefix}-mlp-routed_experts-wo": [
+                    f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices
+                ],
+                (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [
+                    f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices
+                ],
+            }
+        )
   else:
     # Unscanned layer mapping
     for i in range(num_main_layers):
@@ -1013,6 +1094,23 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
                 f"{prefix}-attention-attention-key_norm-scale": f"model.language_model.layers.{i}.self_attn.k_norm.weight",
             }
         )
+        if is_quantized:
+          mapping.update(
+              {
+                  f"{prefix}-attention-attention-query-kernel_scale": (
+                      f"model.language_model.layers.{i}.self_attn.q_proj.weight_scale_inv"
+                  ),
+                  f"{prefix}-attention-attention-key-kernel_scale": (
+                      f"model.language_model.layers.{i}.self_attn.k_proj.weight_scale_inv"
+                  ),
+                  f"{prefix}-attention-attention-value-kernel_scale": (
+                      f"model.language_model.layers.{i}.self_attn.v_proj.weight_scale_inv"
+                  ),
+                  f"{prefix}-attention-attention-out-kernel_scale": (
+                      f"model.language_model.layers.{i}.self_attn.o_proj.weight_scale_inv"
+                  ),
+              }
+          )
       else:
         # Linear/Hybrid Attention Block (Unscanned)
         mapping.update(  # pyrefly: ignore[no-matching-overload]
@@ -1034,6 +1132,18 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
                 f"{prefix}-attention-out_proj-kernel": f"model.language_model.layers.{i}.linear_attn.out_proj.weight",
             }
         )
+        if is_quantized:
+          mapping.update(
+              {
+                  f"{prefix}-attention-in_proj_qkvz-kernel_scale": (
+                      f"model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight_scale_inv",
+                      f"model.language_model.layers.{i}.linear_attn.in_proj_z.weight_scale_inv",
+                  ),
+                  f"{prefix}-attention-out_proj-kernel_scale": (
+                      f"model.language_model.layers.{i}.linear_attn.out_proj.weight_scale_inv"
+                  ),
+              }
+          )
 
       # MLP: Gates and Shared Experts
       hf_mlp = f"model.language_model.layers.{i}.mlp"
@@ -1047,17 +1157,53 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
               f"{prefix}-mlp-shared_expert_gate-kernel": (f"{hf_mlp}.shared_expert_gate.weight"),
           }
       )
+      if is_quantized:
+        mapping.update(
+            {
+                f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": (f"{hf_mlp}.shared_expert.gate_proj.weight_scale_inv"),
+                f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": (f"{hf_mlp}.shared_expert.up_proj.weight_scale_inv"),
+                f"{prefix}-mlp-shared_expert-wo-kernel_scale": (f"{hf_mlp}.shared_expert.down_proj.weight_scale_inv"),
+            }
+        )
 
       # MoE Routed Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj",
-              (
-                  f"{prefix}-mlp-routed_experts-wi_0",
-                  f"{prefix}-mlp-routed_experts-wi_1",
-              ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj",
-          }
-      )
+      if is_quantized:
+        num_experts = config.get("text_config", config).get("num_experts", 256)
+        mapping.update(
+            {
+                f"{prefix}-mlp-routed_experts-wi_0": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_0_scale": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv"
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_1": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wi_1_scale": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv"
+                    for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wo": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts)
+                ],
+                f"{prefix}-mlp-routed_experts-wo_scale": [
+                    f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv"
+                    for e in range(num_experts)
+                ],
+            }
+        )
+      else:
+        mapping.update(  # pyrefly: ignore[no-matching-overload]
+            {
+                f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj",
+                (
+                    f"{prefix}-mlp-routed_experts-wi_0",
+                    f"{prefix}-mlp-routed_experts-wi_1",
+                ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj",
+            }
+        )
 
   # Vision mapping for Qwen3.5
   if maxtext_config.use_multimodal and "vision_config" in config:
@@ -1232,6 +1378,44 @@ def concat_qkvz_and_transpose(input_tensor, target_shape=None):
       interleaved = np.concatenate([q_r, k_r, v_r, z_r], axis=1)
       return interleaved.reshape(-1, qkv_m.shape[-1]).T
 
+  raw_block_size = getattr(maxtext_config, "weight_block_size", None)
+  if raw_block_size is None and isinstance(config, dict):
+    raw_block_size = config.get("quantization_config", {}).get("weight_block_size", 128)
+  weight_block_size = raw_block_size[0] if isinstance(raw_block_size, (list, tuple)) else (raw_block_size or 128)
+
+  def concat_qkvz_scales_and_transpose(input_tensor, target_shape=None):
+    if saving_to_hf:
+      t_m = input_tensor.T
+      t_r = t_m.reshape(H_k, -1, t_m.shape[-1])
+      d_k_blocks = max(1, D_k // weight_block_size)
+      d_v_blocks = max(1, (V_per_K * D_v) // weight_block_size)
+      q_scale = t_r[:, :d_k_blocks, :].reshape(H_k * d_k_blocks, -1)
+      k_scale = t_r[:, d_k_blocks : 2 * d_k_blocks, :].reshape(H_k * d_k_blocks, -1)
+      v_scale = t_r[:, 2 * d_k_blocks : 2 * d_k_blocks + d_v_blocks, :].reshape(
+          H_v * max(1, D_v // weight_block_size), -1
+      )
+      z_scale = t_r[:, 2 * d_k_blocks + d_v_blocks :, :].reshape(H_v * max(1, D_v // weight_block_size), -1)
+      qkv_scale = np.concatenate([q_scale, k_scale, v_scale], axis=0)
+      return qkv_scale, z_scale
+    else:
+      qkv_scale, z_scale = input_tensor
+      d_k_blocks = max(1, D_k // weight_block_size)
+      d_v_blocks = max(1, (V_per_K * D_v) // weight_block_size)
+      Q_blocks = H_k * d_k_blocks
+      K_blocks = H_k * d_k_blocks
+
+      q_scale = qkv_scale[:Q_blocks, :]
+      k_scale = qkv_scale[Q_blocks : Q_blocks + K_blocks, :]
+      v_scale = qkv_scale[Q_blocks + K_blocks :, :]
+
+      q_scale_r = q_scale.reshape(H_k, d_k_blocks, -1)
+      k_scale_r = k_scale.reshape(H_k, d_k_blocks, -1)
+      v_scale_r = v_scale.reshape(H_k, d_v_blocks, -1)
+      z_scale_r = z_scale.reshape(H_k, d_v_blocks, -1)
+
+      interleaved = np.concatenate([q_scale_r, k_scale_r, v_scale_r, z_scale_r], axis=1)
+      return interleaved.reshape(-1, qkv_scale.shape[-1]).T
+
   def concat_ba_and_transpose(input_tensor, target_shape=None):
     if saving_to_hf:
       t_m = input_tensor.T
@@ -1251,6 +1435,11 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
       interleaved = np.concatenate([b_r, a_r], axis=1)
       return interleaved.reshape(-1, b_m.shape[-1]).T
 
+  is_quantized = getattr(maxtext_config, "weight_dtype", None) in (
+      "float8_e4m3fn",
+      "float8_e5m2",
+  )
+
   # Initialize Hooks
   hooks = {
       "params-decoder-logits_dense-kernel": transpose,
@@ -1272,11 +1461,16 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
     if is_full_attention_layer:
       for key in ["query", "key", "value", "out"]:
         hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel  # pyrefly: ignore[bad-assignment]
+        if is_quantized:
+          hooks[f"{prefix}-attention-attention-{key}-kernel_scale"] = reshape_kernel
     else:
       hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = concat_qkvz_and_transpose
       hooks[f"{prefix}-attention-in_proj_ba-kernel"] = concat_ba_and_transpose
       hooks[f"{prefix}-attention-out_proj-kernel"] = transpose
       hooks[f"{prefix}-attention-conv1d-kernel"] = permute_conv
+      if is_quantized:
+        hooks[f"{prefix}-attention-in_proj_qkvz-kernel_scale"] = concat_qkvz_scales_and_transpose
+        hooks[f"{prefix}-attention-out_proj-kernel_scale"] = transpose
 
     mlp_prefix = f"{prefix}-mlp"
     hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose
@@ -1284,11 +1478,22 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
     hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose
     hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose
     hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose
-    # pyrefly: ignore[unsupported-operation]
-    hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = (
-        process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
-    )
-    hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
+    if is_quantized:
+      hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel_scale"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel_scale"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert-wo-kernel_scale"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wi_0"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wi_0_scale"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wi_1_scale"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose
+      hooks[f"{mlp_prefix}-routed_experts-wo_scale"] = transpose
+    else:
+      # pyrefly: ignore[unsupported-operation]
+      hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = (
+          process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
+      )
+      hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
 
   # Vision hooks for Qwen3.5
   vision_config = config.get("vision_config", None)
@@ -4262,7 +4467,10 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3.5-397b-a17b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "qwen3.5-397b-a17b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "qwen3.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "qwen3.5-35b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
     "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "olmo3-7b": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING,
@@ -4316,7 +4524,10 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN,
     "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3.5-397b-a17b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "qwen3.5-397b-a17b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "qwen3.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "qwen3.5-35b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
diff --git a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py
index 508697624b..0fb262faf4 100644
--- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py
+++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py
@@ -122,7 +122,7 @@ def _build_single_axis_stacked_tensor(
 
   if config.scan_layers:
     # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0
   else:
     # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0).
     axis_to_stack = 0
diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py
index 26677a0cbb..252ae617fb 100644
--- a/src/maxtext/checkpoint_conversion/utils/utils.py
+++ b/src/maxtext/checkpoint_conversion/utils/utils.py
@@ -287,7 +287,12 @@ def process_maxtext_param(
     if maxtext_config.scan_layers:
       max_logging.log("\tscan")
       # Case 2: Standard scanned layer. Stacked ONLY on the layer axis.
-      axis_to_slice = maxtext_config.param_scan_axis
+      weight_sample = maxtext_param_weight[0] if isinstance(maxtext_param_weight, list) else maxtext_param_weight
+      axis_to_slice = (
+          maxtext_config.param_scan_axis
+          if getattr(weight_sample, "ndim", 0) > maxtext_config.param_scan_axis
+          else 0
+      )
     else:
       max_logging.log("\tunscan moe")
       # Case 3: Unscanned MoE layer. Stacked ONLY on the expert axis. Assuming expert is axis 0.
@@ -1373,7 +1378,7 @@ def _build_single_axis_stacked_tensor(
 
   if config.scan_layers:
     # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0
   else:
     # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0).
     axis_to_stack = 0
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index 7191190f2b..21c0f46d70 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -30,10 +30,12 @@
 from flax.training import train_state
 from grain.experimental import ElasticIterator
 import jax
+import jax.numpy as jnp
 from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state
 from maxtext.common import emergency_checkpointing
 from maxtext.common import grain_utility
 from maxtext.common import train_state_nnx
+from maxtext.layers import linears
 from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator
 from maxtext.input_pipeline.multihost_dataloading import RemoteIteratorWrapper
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
@@ -701,6 +703,142 @@ def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-
   return orbax_cloud_logger
 
 
+def _get_checkpoint_metadata_tree(ckptr: Any, ckpt_path: epath.Path) -> Any:
+  """Safely retrieves the metadata tree for an Orbax checkpoint."""
+  try:
+    metadata = ckptr.metadata(ckpt_path)
+    if metadata is None:
+      return None
+    if hasattr(metadata, "item_metadata"):
+      metadata = metadata.item_metadata
+    if hasattr(metadata, "tree"):
+      return metadata.tree
+    return metadata
+  except Exception as e:  # pylint: disable=broad-except
+    max_logging.log(f"Warning: Failed to retrieve checkpoint metadata from {ckpt_path}: {e}")
+    return None
+
+
+def _find_matching_meta_subtree(want_bare: Any, meta_tree: Any) -> Any:
+  """Unwraps metadata tree wrappers (e.g. 'params', 'model_params') to align with want_bare."""
+  if not isinstance(want_bare, dict) or not isinstance(meta_tree, dict):
+    return meta_tree
+
+  want_keys = set(want_bare.keys())
+  if want_keys and want_keys.issubset(meta_tree.keys()):
+    return meta_tree
+
+  for wrapper in ("params", "model_params", "model", "items"):
+    if wrapper in meta_tree and isinstance(meta_tree[wrapper], dict):
+      sub = meta_tree[wrapper]
+      if want_keys and want_keys.issubset(sub.keys()):
+        return sub
+      if wrapper == "params" and "params" in sub and isinstance(sub["params"], dict):
+        if want_keys and want_keys.issubset(sub["params"].keys()):
+          return sub["params"]
+
+  return meta_tree
+
+
+def _augment_target_with_scales(want_node: Any, meta_node: Any) -> Any:
+  """Augments `want_node` with companion scales from `meta_node` if present in checkpoint but not in want."""
+  if not isinstance(want_node, dict) or not isinstance(meta_node, dict):
+    return want_node
+
+  augmented = {}
+  for k, v in want_node.items():
+    scale_key = f"{k}_scale"
+    meta_scale = meta_node.get(scale_key) if isinstance(meta_node, dict) else None
+    if not isinstance(v, dict) and scale_key not in want_node and meta_scale is not None:
+      # Dequantize-on-load: Target wants unquantized weights, but checkpoint has companion scale.
+      meta_param = meta_node.get(k)
+      param_dtype = getattr(meta_param, "dtype", getattr(v, "dtype", jnp.bfloat16))
+      augmented[k] = jax.ShapeDtypeStruct(
+          shape=getattr(v, "shape", getattr(meta_param, "shape", ())),
+          dtype=param_dtype,
+          sharding=getattr(v, "sharding", None),
+      )
+      augmented[scale_key] = jax.ShapeDtypeStruct(
+          shape=getattr(meta_scale, "shape", ()),
+          dtype=getattr(meta_scale, "dtype", jnp.float32),
+          sharding=None,
+      )
+    elif isinstance(v, dict):
+      augmented[k] = _augment_target_with_scales(
+          v,
+          meta_node.get(k) if isinstance(meta_node, dict) else None,
+      )
+    else:
+      if k not in augmented:
+        augmented[k] = v
+
+  return augmented
+
+
+def _augment_want_with_scales(want: Any, meta_tree: Any, is_nnx: bool, restore_key: str) -> Any:
+  """Augments the target params dictionary with scales from checkpoint metadata if needed."""
+  if meta_tree is None:
+    return want
+
+  if is_nnx or restore_key in ("model_params", "model"):
+    meta_weights = _find_matching_meta_subtree(want, meta_tree)
+    return _augment_target_with_scales(want, meta_weights)
+  else:
+    # Linen: want is {"params": bare_weights} or bare_weights
+    if isinstance(want, dict) and "params" in want and len(want) == 1:
+      want_bare = want["params"]
+      meta_weights = _find_matching_meta_subtree(want_bare, meta_tree)
+      augmented_bare = _augment_target_with_scales(want_bare, meta_weights)
+      return {"params": augmented_bare}
+    else:
+      meta_weights = _find_matching_meta_subtree(want, meta_tree)
+      return _augment_target_with_scales(want, meta_weights)
+
+
+def maybe_dequantize_restored_params(restored_weights: Any, want: Any) -> Any:
+  """Dequantizes restored weights if checkpoint contained companion scales but want did not.
+
+  For each parameter dictionary containing a weight and its companion scale (e.g. 'kernel' and
+  'kernel_scale', or 'wi_0' and 'wi_0_scale') where the target model (`want`) does not expect the scale,
+  dynamically dequantizes the weight to the target compute dtype using `dequantize_weight` from
+  `maxtext.layers.linears` and removes the companion scale from the restored parameter dictionary.
+
+  Args:
+    restored_weights: The restored parameter PyTree from the checkpoint.
+    want: The expected parameter PyTree structure / ShapeDtypeStructs.
+
+  Returns:
+    The parameter PyTree with dequantized weights and omitted companion scales where applicable.
+  """
+  if not isinstance(restored_weights, dict):
+    return restored_weights
+
+  consumed_scales = set()
+  out = {}
+
+  for k, v in restored_weights.items():
+    scale_key = f"{k}_scale"
+    has_scale = scale_key in restored_weights
+    want_expects_scale = isinstance(want, dict) and scale_key in want
+
+    if has_scale and not want_expects_scale:
+      target_param = want.get(k) if isinstance(want, dict) else None
+      target_dtype = getattr(target_param, "dtype", jnp.bfloat16)
+
+      weight = v
+      scale = restored_weights[scale_key]
+      out[k] = linears.dequantize_weight(weight, scale, compute_dtype=target_dtype)
+      consumed_scales.add(scale_key)
+
+  for k, v in restored_weights.items():
+    if k in consumed_scales or k in out:
+      continue
+    want_sub = want.get(k) if isinstance(want, dict) else None
+    out[k] = maybe_dequantize_restored_params(v, want_sub)
+
+  return out
+
+
 def load_params_from_path(
     load_parameters_from_path,
     abstract_unboxed_params,
@@ -723,11 +861,6 @@ def load_params_from_path(
   if restore_key not in ("model_params", "model"):
     restore_key = "params"
 
-  if restore_key in ("model_params", "model"):
-    params_collection = want
-  else:
-    params_collection = {"params": want} if is_nnx else want
-
   # *_concurrent_gb should be set for large models, the default is 96.
   max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}")
   ckptr = ocp.Checkpointer(
@@ -739,13 +872,22 @@ def load_params_from_path(
       )
   )
 
+  ckpt_path = epath.Path(load_parameters_from_path)
+  meta_tree = _get_checkpoint_metadata_tree(ckptr, ckpt_path)
+  augmented_want = _augment_want_with_scales(want, meta_tree, is_nnx, restore_key)
+
+  if restore_key in ("model_params", "model"):
+    params_collection = augmented_want
+  else:
+    params_collection = {"params": augmented_want} if is_nnx else augmented_want
+
   # This is a memory optimization. We don't want to restore the entire checkpoint - only the params.
   # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste
   # memory, we instead specify here that we are just restoring the params field of the checkpoint
   # (which itself may be a dictionary containing a key named 'params' or 'model').
   restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection)
   restored = ckptr.restore(
-      epath.Path(load_parameters_from_path),
+      ckpt_path,
       item={restore_key: params_collection},
       transforms={},
       restore_args={restore_key: restore_args},
@@ -757,6 +899,11 @@ def load_params_from_path(
   else:
     restored_weights = restored_collection["params"] if is_nnx else restored_collection
 
+  # Dequantize if checkpoint had companion kernel_scale and target want is unquantized.
+  restored_weights = maybe_dequantize_restored_params(restored_weights, want)
+  if not is_nnx:
+    restored_collection = restored_weights
+
   # `transforms={}` lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks,
   # and a stored array at its own shape rather than the target's. Either reaches the model and
   # fails much later without naming the weight, so check here -- the params-only load
diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py
index 0de9be9085..6353db876f 100644
--- a/src/maxtext/common/common_types.py
+++ b/src/maxtext/common/common_types.py
@@ -14,6 +14,7 @@
 
 """Common types."""
 import enum
+import fnmatch
 from typing import Any, Sequence
 
 import numpy as np
@@ -28,6 +29,29 @@
 DType = jnp.dtype
 Shape = Sequence[int]
 
+
+def is_fp8_dtype(dtype: Any) -> bool:
+  """Checks whether a dtype is FP8."""
+  return dtype in (
+      "float8_e4m3fn",
+      "float8_e5m2",
+      jnp.float8_e4m3fn,
+      jnp.float8_e5m2,
+  )
+
+
+def get_weight_dtype(config: Config, module_name: str) -> DType:
+  """Resolves parameter storage dtype for a submodule, honoring unquantized_modules."""
+  if not is_fp8_dtype(config.weight_dtype):
+    return config.weight_dtype
+  unquantized = getattr(config, "unquantized_modules", None) or ()
+  leaf_name = module_name.rsplit(".", 1)[-1]
+  if any(fnmatch.fnmatch(module_name, p) or fnmatch.fnmatch(leaf_name, p) for p in unquantized):
+    return config.dtype
+  return config.weight_dtype
+
+
+
 AxisNames = tuple[str, ...]
 AxisIdxes = tuple[int, ...]
 
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index f3341ce211..480052c547 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -172,6 +172,10 @@ decoder_block: "llama2" # which style of decoderblock to use.
 # then you should explicitly set base_embed_dim, base_num_query_heads, base_num_kv_heads,
 # base_mlp_dim, base_num_decoder_layers and/or head_dim.
 weight_dtype: "float32"
+# List of module names or patterns to keep unquantized in model dtype (e.g. bfloat16) even when weight_dtype is FP8
+unquantized_modules: []
+# Block size for block-scaled quantized weights (e.g. 128 for 128x128 block scaling). null for per-tensor scaling.
+weight_block_size: null
 global_parameter_scale: 1
 base_emb_dim: 2048
 base_num_query_heads: 16
diff --git a/src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml b/src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml
new file mode 100644
index 0000000000..e8f597ed05
--- /dev/null
+++ b/src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml
@@ -0,0 +1,93 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# model config for qwen3.5-35b-a3b-fp8 (FP8 weight-only storage with dynamic dequantization)
+
+base_config: "base.yml"
+
+decoder_block: "qwen3_5"
+
+# Core Architectural Parameters
+base_emb_dim: 2048
+base_num_decoder_layers: 40
+base_num_query_heads: 16
+base_num_kv_heads: 2
+head_dim: 256
+vocab_size: 248320
+normalization_layer_epsilon: 1.0e-6
+
+# MoE Specific Parameters
+# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models.
+base_mlp_dim: 512
+base_moe_mlp_dim: 512
+num_experts: 256
+shared_experts: 1
+num_experts_per_tok: 8
+norm_topk_prob: True
+float32_gate_logits: True
+sparse_matmul: False
+
+# GatedDeltaNet Specific Parameters for Linear Attention (GDN)
+inhomogeneous_layer_cycle_interval: 4
+gdn_conv_kernel_dim: 4
+gdn_key_head_dim: 128
+gdn_value_head_dim: 128
+gdn_num_key_heads: 16
+gdn_num_value_heads: 32
+gdn_chunk_size: 64
+
+# RoPE Settings
+rope_max_timescale: 10000000
+partial_rotary_factor: 0.25
+
+# General Model Settings
+enable_dropout: False
+logits_via_embedding: False
+
+# Vision Encoder Configuration (need to set use_multimodal=true)
+vision_encoder_block: "qwen3_5"
+# Based on Qwen3.5 MoE Vision Model Config
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 2048  # Projects to decoder emb_dim (2048)
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: []  # No deepstack for Qwen3.5 VL
+rope_theta_for_vit: 10000
+
+# MRoPE Settings (Multi-dimensional RoPE for multimodal)
+use_mrope: true
+mrope_section: [11, 11, 10]
+
+weight_dtype: "float8_e4m3fn"
+weight_block_size: 128
+dtype: "bfloat16"
+unquantized_modules: [
+  "token_embedder",
+  "logits_dense",
+  "gate",
+  "shared_expert_gate",
+  "conv1d",
+  "in_proj_ba",
+  "A_log",
+  "dt_bias",
+  "norm"
+]
+
diff --git a/src/maxtext/configs/models/qwen3.5-397b-a17b-fp8.yml b/src/maxtext/configs/models/qwen3.5-397b-a17b-fp8.yml
new file mode 100644
index 0000000000..85e284ef37
--- /dev/null
+++ b/src/maxtext/configs/models/qwen3.5-397b-a17b-fp8.yml
@@ -0,0 +1,88 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# model config for qwen3.5-397b-a17b-fp8 (FP8 weight-only storage with dynamic dequantization)
+
+decoder_block: "qwen3_5"
+
+# Core Architectural Parameters
+base_emb_dim: 4096
+base_num_decoder_layers: 60
+base_num_query_heads: 32
+base_num_kv_heads: 2
+head_dim: 256
+vocab_size: 248320
+normalization_layer_epsilon: 1.0e-6
+
+# MoE Specific Parameters
+# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models.
+base_mlp_dim: 1024
+base_moe_mlp_dim: 1024
+num_experts: 512
+shared_experts: 1
+num_experts_per_tok: 10
+norm_topk_prob: True
+float32_gate_logits: True
+
+# GatedDeltaNet Specific Parameters for Linear Attention (GDN)
+inhomogeneous_layer_cycle_interval: 4
+gdn_conv_kernel_dim: 4
+gdn_key_head_dim: 128
+gdn_value_head_dim: 128
+gdn_num_key_heads: 16
+gdn_num_value_heads: 64
+gdn_chunk_size: 64
+
+# RoPE Settings
+rope_max_timescale: 10000000
+partial_rotary_factor: 0.25
+
+# General Model Settings
+enable_dropout: False
+
+# Vision Encoder Configuration (need to set use_multimodal=true)
+vision_encoder_block: "qwen3_5"
+# Based on Qwen3.5 MoE Vision Model Config
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 4096  # Projects to decoder emb_dim (4096)
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: []  # No deepstack for Qwen3.5 VL
+rope_theta_for_vit: 10000
+
+# MRoPE Settings (Multi-dimensional RoPE for multimodal)
+use_mrope: true
+mrope_section: [11, 11, 10]
+
+weight_dtype: "float8_e4m3fn"
+weight_block_size: 128
+dtype: "bfloat16"
+unquantized_modules: [
+  "token_embedder",
+  "logits_dense",
+  "gate",
+  "shared_expert_gate",
+  "conv1d",
+  "in_proj_ba",
+  "A_log",
+  "dt_bias",
+  "norm"
+]
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index f771b056fc..9c865159d8 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -66,6 +66,8 @@ class DType(str, Enum):
   BFLOAT16 = "bfloat16"
   FLOAT32 = "float32"
   FLOAT16 = "float16"
+  FLOAT8_E4M3FN = "float8_e4m3fn"
+  FLOAT8_E5M2 = "float8_e5m2"
 
 
 class MatmulPrecision(str, Enum):
@@ -279,7 +281,10 @@ class ProfilerType(str, Enum):
     "qwen3-omni-30b-a3b",
     "qwen3-custom-30b-a3b",
     "qwen3.5-35b-a3b",
+    "qwen3.5-35b-a3b-fp8",
+    "qwen3.5-35b-fp8",
     "qwen3.5-397b-a17b",
+    "qwen3.5-397b-a17b-fp8",
     "gpt3-175b",
     "gpt3-22b",
     "gpt3-6b",
@@ -457,6 +462,20 @@ class Quantization(BaseModel):
       QuantizationType.NONE,
       description="Activates quantization for transformer layers.",
   )
+  unquantized_modules: list[str] = Field(
+      default_factory=list,
+      description=(
+          "List of submodule names or name patterns to keep unquantized in model dtype "
+          "(e.g. bfloat16) even when weight_dtype is FP8."
+      ),
+  )
+  weight_block_size: None | int | list[int] = Field(
+      None,
+      description=(
+          "Block size for block-scaled quantized weights (e.g. 128 for 128x128 block scaling). "
+          "None for per-tensor scaling."
+      ),
+  )
   replicate_quant_scale: bool = Field(
       False,
       description="Replicates quantization scale to avoid inefficient XLA fusion.",
diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py
index c61b945b0e..962aa54493 100644
--- a/src/maxtext/layers/attentions.py
+++ b/src/maxtext/layers/attentions.py
@@ -685,6 +685,7 @@ def query_init(*args):
     if self.is_qwen3_hybrid:
       out_features = (self.num_query_heads, self.head_dim * 2)
 
+    block_size = getattr(self.config, "weight_block_size", None)
     return DenseGeneral(
         in_features_shape=in_features,
         out_features_shape=out_features,
@@ -697,6 +698,7 @@ def query_init(*args):
         matmul_precision=self.config.matmul_precision,
         use_bias=self.use_bias_in_projections,
         shard_mode=self.config.shard_mode,
+        block_size=block_size,
         rngs=self.rngs,
     )
 
@@ -723,6 +725,7 @@ def init_kv_w(self, inputs_kv_shape: Tuple) -> nnx.Module:
     )
     self._validate_kv_head_sharding(kernel_axes)
 
+    block_size = getattr(self.config, "weight_block_size", None)
     return DenseGeneral(
         in_features_shape=self.convert_dense_general_inputs_shape(inputs_kv_shape),
         out_features_shape=(self.num_kv_heads, self.head_dim),
@@ -735,6 +738,7 @@ def init_kv_w(self, inputs_kv_shape: Tuple) -> nnx.Module:
         shard_mode=self.config.shard_mode,
         matmul_precision=self.config.matmul_precision,
         use_bias=self.use_bias_in_projections,
+        block_size=block_size,
         rngs=self.rngs,
     )
 
@@ -819,6 +823,7 @@ def init_out_w(self, output_dim: int) -> nnx.Module:
       out_kernel_axis = ("mlp", "embed_attn")
       axis = (-1,)
 
+    block_size = getattr(self.config, "weight_block_size", None)
     return DenseGeneral(
         in_features_shape=in_features,
         out_features_shape=out_features,
@@ -831,6 +836,7 @@ def init_out_w(self, output_dim: int) -> nnx.Module:
         shard_mode=self.config.shard_mode,
         matmul_precision=self.config.matmul_precision,
         use_bias=False if self.is_qwen2 else self.use_bias_in_projections,
+        block_size=block_size,
         rngs=self.rngs,
     )
 
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index 18a9906c15..5844f96e21 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -27,7 +27,7 @@
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from maxtext.common.common_types import Config, DecoderBlockType, ShardMode
+from maxtext.common.common_types import Config, DecoderBlockType, ShardMode, get_weight_dtype
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
 from maxtext.layers import linears
 from maxtext.layers import mhc
@@ -722,6 +722,8 @@ def _apply_embedding(
             "qwen3-vl-4b",
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
+            "qwen3.5-35b-a3b-fp8",
+            "qwen3.5-35b-fp8",
             "qwen3.5-397b-a17b",
             "maxtext-omni-gemma3-qwen3",
         ]:
@@ -742,6 +744,8 @@ def _apply_embedding(
             "qwen3-vl-4b",
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
+            "qwen3.5-35b-a3b-fp8",
+            "qwen3.5-35b-fp8",
             "qwen3.5-397b-a17b",
         ]:
           y = mm_utils.merge_mm_embeddings(
@@ -828,10 +832,11 @@ def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, determi
         logits = logits / cfg.final_logits_soft_cap
         logits = jnp.tanh(logits) * cfg.final_logits_soft_cap
     else:
+      logits_weight_dtype = get_weight_dtype(cfg, "logits_dense")
       logits = linears.dense_general(
           inputs_shape=y.shape,
           out_features_shape=cfg.vocab_size,
-          weight_dtype=cfg.weight_dtype,
+          weight_dtype=logits_weight_dtype,
           dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype,  # for logit training stability
           kernel_axes=("embed_vocab", "vocab"),
           shard_mode=cfg.shard_mode,
diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py
index a744ea041a..fb8a0441ba 100644
--- a/src/maxtext/layers/embeddings.py
+++ b/src/maxtext/layers/embeddings.py
@@ -25,7 +25,7 @@
 
 from flax import nnx
 
-from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType
+from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType, get_weight_dtype
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned
 from maxtext.utils import max_logging
@@ -126,12 +126,12 @@ def __init__(
     self.cast_input_dtype = cast_input_dtype
     self.dtype = dtype
     self.attend_dtype = attend_dtype
-
+    embed_weight_dtype = get_weight_dtype(self.config, "token_embedder")
     self.embedding = nnx.Param(
         embedding_init(
             rngs.params(),
             (self.num_embeddings, self.num_features),
-            self.config.weight_dtype,
+            embed_weight_dtype,
         ),
         sharding=("vocab", "embed_vocab"),
     )
diff --git a/src/maxtext/layers/initializers.py b/src/maxtext/layers/initializers.py
index bbc6605057..742e00dab5 100644
--- a/src/maxtext/layers/initializers.py
+++ b/src/maxtext/layers/initializers.py
@@ -17,18 +17,27 @@
 from typing import Callable
 
 import jax
+import jax.numpy as jnp
 
 from flax import linen as nn
 from flax import nnx
 from aqt.jax.v2 import aqt_tensor
 
-from maxtext.common.common_types import Array, DType, Shape, PRNGKey
+from maxtext.common.common_types import Array, DType, Shape, PRNGKey, is_fp8_dtype
 
 Initializer = Callable[[PRNGKey, Shape, DType], Array]
 InitializerAxis = int | tuple[int, ...]
 NdInitializer = Callable[[PRNGKey, Shape, DType, InitializerAxis, InitializerAxis], Array]
 
-default_embed_init = nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0)
+
+def _default_embed_init(key, shape, dtype=jnp.float32):
+  target_dtype = dtype
+  sample_dtype = jnp.float32 if is_fp8_dtype(dtype) else dtype
+  fn = nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0)
+  return fn(key, shape, sample_dtype).astype(target_dtype)
+
+
+default_embed_init = _default_embed_init
 
 default_bias_init = jax.nn.initializers.constant(0.0)
 default_scalar_init = jax.nn.initializers.constant(0.01)
@@ -54,8 +63,10 @@ def nd_dense_init(scale, mode, distribution):
 
   def init_fn(key, shape, dtype, in_axis, out_axis):
     """Initializes an array using variance scaling with specified axes."""
+    target_dtype = dtype
+    sample_dtype = jnp.float32 if is_fp8_dtype(dtype) else dtype
     fn = jax.nn.initializers.variance_scaling(scale, mode, distribution, in_axis, out_axis)
-    return fn(key, shape, dtype)
+    return fn(key, shape, sample_dtype).astype(target_dtype)
 
   return init_fn
 
diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py
index e088f7ec50..1f5be38227 100644
--- a/src/maxtext/layers/linears.py
+++ b/src/maxtext/layers/linears.py
@@ -29,11 +29,11 @@
 from flax import nnx
 import flax.linen as nn
 
-from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config
+from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config, Shape, is_fp8_dtype
 from maxtext.common.common_types import MODEL_MODE_PREFILL
 from maxtext.layers import nnx_wrappers, quantizations
 from maxtext.layers import normalizations
-from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned
+from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned, Initializer
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
@@ -74,7 +74,54 @@ def canonicalize_tuple(x):
     return (x,)
 
 
-def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul_precision, quant):
+def dequantize_weight(
+    w: Array,
+    scale: Array | None,
+    compute_dtype: DType = jnp.bfloat16,
+) -> Array:
+  """Dequantizes weight tensor `w` dynamically to `compute_dtype` using `scale`.
+
+  Supports:
+    1. No scale (scale is None): casts w to compute_dtype.
+    2. Per-tensor scalar scale: scalar broadcast multiplication.
+    3. Block-wise scale: 2D/3D block broadcast multiplication.
+    4. General broadcastable scale: standard JAX broadcasting.
+  """
+  if scale is None:
+    return w.astype(compute_dtype)
+
+  w_c = w.astype(compute_dtype)
+  scale_c = jnp.asarray(scale, compute_dtype)
+
+  # Per-tensor scalar scale or matching shape
+  if scale_c.ndim == 0 or scale_c.shape == w.shape:
+    return w_c * scale_c
+
+  # Block-wise scale (e.g. 2D for Dense, 3D for MoE)
+  if scale_c.ndim == w.ndim and any(s > 1 and s != d for s, d in zip(scale_c.shape, w.shape)):
+    if not all(d % s == 0 for d, s in zip(w.shape, scale_c.shape)):
+      raise ValueError(
+          f"Block scaling requires weight dimensions {w.shape} to be divisible by scale dimensions {scale_c.shape}."
+      )
+    interleaved_shape = tuple(dim for s, w_d in zip(scale_c.shape, w.shape) for dim in (s, w_d // s))
+    scale_shape = tuple(dim for s in scale_c.shape for dim in (s, 1))
+    return (w_c.reshape(interleaved_shape) * scale_c.reshape(scale_shape)).reshape(w.shape)
+
+  # Standard JAX broadcasting handles per-channel or broadcastable shapes
+  return w_c * scale_c
+
+
+def _compute_dot_general(
+    inputs,
+    kernel,
+    kernel_axes,
+    axis,
+    contract_ind,
+    matmul_precision,
+    quant,
+    kernel_scale: Array | None = None,
+    compute_dtype: DType | None = None,
+):
   """Computes a dot_general operation that may be quantized."""
   dot_general = lax.dot_general
   matmul_precision = lax.Precision(matmul_precision)
@@ -82,6 +129,15 @@ def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul
     dot_general_cls = quant.dot_general_cls(mesh_axes=kernel_axes)
     dot_general = dot_general_cls()
     return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=None)
+
+  if compute_dtype is None:
+    compute_dtype = inputs.dtype
+
+  if is_fp8_dtype(kernel.dtype) or kernel_scale is not None:
+    kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=compute_dtype)
+  else:
+    kernel = jnp.asarray(kernel, compute_dtype)
+
   return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision)
 
 
@@ -94,6 +150,8 @@ def _compute_dot_general_nnx(
     quant_dot_general: nnx_wrappers.ToNNX | None,
     initializing: bool,
     out_sharding: NamedSharding | None = None,
+    kernel_scale: Array | None = None,
+    compute_dtype: DType | None = None,
 ):
   """Computes a dot_general operation that may be quantized."""
   dot_general = lax.dot_general
@@ -107,6 +165,14 @@ def _compute_dot_general_nnx(
     out_ndim = (inputs.ndim - len(axis)) + (kernel.ndim - len(contract_ind))
     out_sharding = truncate_out_sharding(out_sharding, out_ndim)
 
+  if compute_dtype is None:
+    compute_dtype = inputs.dtype
+
+  if is_fp8_dtype(kernel.dtype) or kernel_scale is not None:
+    kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=compute_dtype)
+  else:
+    kernel = jnp.asarray(kernel, compute_dtype)
+
   return dot_general(
       inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision, out_sharding=out_sharding
   )
@@ -132,6 +198,12 @@ def __init__(
       mesh: Mesh | None = None,
       use_two_stage_all_gather: bool = False,
       debug_sharding: bool = False,
+      has_scale: bool | None = None,
+      kernel_scale_init: Initializer | None = None,
+      scale_shape: Shape | None = None,
+      scale_axes: tuple[None | str, ...] | None = None,
+      scale_dtype: DType = jnp.float32,
+      block_size: int | tuple[int, ...] | None = None,
       *,  # Following arguments are keyword-only
       rngs: nnx.Rngs = None,
   ):
@@ -158,6 +230,12 @@ def __init__(
         transpose XLA emits for a single combined 2-axis all-gather.
       debug_sharding: when True, log the logical/physical sharding of the
         two-stage all-gather constraints to the sharding dump files.
+      has_scale: whether to initialize a separate scale parameter (kernel_scale).
+      kernel_scale_init: initializer function for kernel_scale.
+      scale_shape: explicit shape for kernel_scale.
+      scale_axes: logical axes for partitioning kernel_scale.
+      scale_dtype: dtype of kernel_scale (default: float32).
+      block_size: block size for block-wise quantization scales.
       rngs: RNG state for initialization in nnx.
     """
     self.in_features_shape = canonicalize_tuple(in_features_shape)
@@ -175,6 +253,9 @@ def __init__(
     self.mesh = mesh
     self.use_two_stage_all_gather = use_two_stage_all_gather
     self.debug_sharding = debug_sharding
+    self.has_scale = has_scale
+    self.scale_dtype = scale_dtype
+    self.block_size = block_size
 
     # Parameter initialization
     kernel_shape = self.in_features_shape + self.out_features_shape
@@ -182,27 +263,91 @@ def __init__(
     kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features_shape))
 
     if not quantizations.in_serve_mode(self.quant):
+      init_dtype = jnp.float32 if is_fp8_dtype(self.weight_dtype) else self.weight_dtype
+      kernel_val = self.kernel_init(
+          rngs.params(),
+          kernel_shape,
+          init_dtype,
+          kernel_in_axis,
+          kernel_out_axis,
+      ).astype(self.weight_dtype)
+
       self.kernel = nnx.Param(
-          self.kernel_init(
-              rngs.params(),
-              kernel_shape,
-              self.weight_dtype,
-              kernel_in_axis,
-              kernel_out_axis,
-          ),
+          kernel_val,
           sharding=self.kernel_axes,
       )
 
     if self.use_bias:
       bias_axes = self.kernel_axes[-len(self.out_features_shape) :]
       bias_shape = kernel_shape[-len(self.out_features_shape) :]
+      bias_val = default_bias_init(rngs.params(), bias_shape, self.weight_dtype)
       self.bias = nnx.Param(
-          default_bias_init(rngs.params(), bias_shape, self.weight_dtype),
+          bias_val,
           sharding=bias_axes,
       )
     else:
       self.bias = None
 
+    should_have_scale = is_fp8_dtype(self.weight_dtype) if has_scale is None else has_scale
+    if should_have_scale and not quantizations.in_serve_mode(self.quant):
+      # Phase 1: Resolve scale shape based on quantization granularity
+      # - Explicit scale_shape: user or caller override.
+      # - Block scaling (e.g. block_size=128): scales are partitioned into a grid
+      #   of size (K // 128, N // 128). Dimensions smaller than block_size (e.g. head_dim < 128)
+      #   are preserved as-is.
+      # - Per-tensor scaling: a single scalar float32 scale with empty shape ().
+      if scale_shape is not None:
+        resolved_scale_shape = canonicalize_tuple(scale_shape)
+      elif block_size is not None:
+        if isinstance(block_size, int):
+          block_sizes = (block_size,) * len(kernel_shape)
+        elif len(block_size) == len(kernel_shape):
+          block_sizes = tuple(block_size)
+        else:
+          block_sizes = (block_size[0],) * len(kernel_shape)
+        resolved_scale_shape = tuple(d if d < b else d // b for d, b in zip(kernel_shape, block_sizes))
+      else:
+        resolved_scale_shape = ()
+
+      # Phase 2: Resolve scale sharding axes to match the weight tensor's mesh partitioning
+      # - Scalar scale (): cannot be sharded across mesh devices, so sharding is empty ().
+      # - A scale dimension whose size equals the kernel's real dimension size (e.g. an
+      #   explicit per-channel scale) is exactly as shardable as the kernel itself, so it
+      #   inherits kernel_axes.
+      # - A scale dimension that's smaller than the kernel's real dimension (block-compressed
+      #   by block_size, e.g. embed_dim=4096 -> 32 blocks, or collapsed to 1) is always far
+      #   smaller than what the kernel's mesh axis is sized for. Inheriting kernel_axes there
+      #   would shard this tiny array along the same physical mesh axis as the (much larger)
+      #   kernel, which silently breaks once FSDP/TP degree exceeds the compressed dimension
+      #   size (e.g. 256 FSDP ways on a 32-block axis). Replicate instead: the array is cheap
+      #   enough that replication costs nothing and is always divisibility-safe.
+      if scale_axes is not None:
+        resolved_scale_axes = scale_axes
+      elif len(resolved_scale_shape) == 0:
+        resolved_scale_axes = ()
+      elif len(resolved_scale_shape) == len(kernel_shape):
+        padded_kernel_axes = self.kernel_axes + (None,) * (len(kernel_shape) - len(self.kernel_axes))
+        resolved_scale_axes = tuple(
+            ax if s_dim == k_dim else None
+            for ax, s_dim, k_dim in zip(padded_kernel_axes, resolved_scale_shape, kernel_shape)
+        )
+      else:
+        resolved_scale_axes = tuple(None for _ in resolved_scale_shape)
+
+      actual_scale_init = kernel_scale_init if kernel_scale_init is not None else jax.nn.initializers.ones
+      self.scale_axes = resolved_scale_axes
+      self.kernel_scale = nnx.Param(
+          actual_scale_init(
+              rngs.params(),
+              resolved_scale_shape,
+              self.scale_dtype,
+          ),
+          sharding=resolved_scale_axes,
+      )
+    else:
+      self.scale_axes = None
+      self.kernel_scale = None
+
     if quant:
       dot_general_cls = quant.dot_general_cls(mesh_axes=kernel_axes)
       dot_general_linen = dot_general_cls()
@@ -286,19 +431,26 @@ def __call__(
     if quantizations.in_serve_mode(self.quant):
       kernel_shape = self.in_features_shape + self.out_features_shape
       kernel = jnp.zeros(kernel_shape, dtype=self.dtype)
+      kernel_scale = None
     else:
-      kernel = getattr(self.kernel, "value", self.kernel)
-      if hasattr(kernel, "value"):
-        kernel = kernel.value
+      kernel = self.kernel[...]
       # Move logit_dense kernel to device if parameter offloading is enabled
       if self.parameter_memory_host_offload:
         max_logging.log("linear.py: Moving parameter logits_dense kernel to device")
         kernel = jax.device_put(kernel, max_utils.device_space())
-      kernel = jnp.asarray(kernel, self.dtype)
+      if self.kernel_scale is not None:
+        kernel_scale = self.kernel_scale[...]
+        if self.parameter_memory_host_offload:
+          kernel_scale = jax.device_put(kernel_scale, max_utils.device_space())
+      else:
+        kernel_scale = None
 
     if slice_bounds is not None:
       if self.quant is not None:
         raise ValueError("sliced contraction is only supported when quant is None")
+      if is_fp8_dtype(kernel.dtype) or kernel_scale is not None:
+        kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=self.dtype)
+        kernel_scale = None
       begin, end = slice_bounds
       if not 0 <= begin < end <= kernel.shape[-1]:
         raise ValueError(f"slice_bounds {slice_bounds} must be valid and within [0, {kernel.shape[-1]}]")
@@ -320,6 +472,8 @@ def __call__(
         self.quant_dot_general if slice_bounds is None else None,
         _initializing,
         out_sharding,
+        kernel_scale=kernel_scale,
+        compute_dtype=self.dtype,
     )
 
     if self.bias is not None:
@@ -346,6 +500,12 @@ def dense_general(
     shard_mode: ShardMode = ShardMode.AUTO,
     matmul_precision: str = "default",
     parameter_memory_host_offload: bool = False,
+    has_scale: bool | None = None,
+    kernel_scale_init: Initializer | None = None,
+    scale_shape: Shape | None = None,
+    scale_axes: tuple[None | str, ...] | None = None,
+    scale_dtype: DType = jnp.float32,
+    block_size: int | tuple[int, ...] | None = None,
     name: None | str = None,
 ):
   """Creates a DenseGeneral Linen module using nnx.bridge.to_linen.
@@ -365,6 +525,12 @@ def dense_general(
     shard_mode: indicating the shard mode
     matmul_precision: Precision for matrix multiplication.
     parameter_memory_host_offload: Determines whether to offload params to host
+    has_scale: whether to initialize a separate scale parameter (kernel_scale).
+    kernel_scale_init: initializer function for kernel_scale.
+    scale_shape: explicit shape for kernel_scale.
+    scale_axes: logical axes for partitioning kernel_scale.
+    scale_dtype: dtype of kernel_scale (default: float32).
+    block_size: block size for block-wise quantization scales.
     name: name passed to the ToLinen Module
   """
   if not (inputs_shape is not None) ^ (in_features_shape is not None):
@@ -389,6 +555,12 @@ def dense_general(
       shard_mode=shard_mode,
       matmul_precision=matmul_precision,
       parameter_memory_host_offload=parameter_memory_host_offload,
+      has_scale=has_scale,
+      kernel_scale_init=kernel_scale_init,
+      scale_shape=scale_shape,
+      scale_axes=scale_axes,
+      scale_dtype=scale_dtype,
+      block_size=block_size,
       name=name,
       metadata_fn=variable_to_logically_partitioned,
       abstract_init=False,
@@ -489,6 +661,8 @@ def __init__(
     else:
       self.intermediate_logical = ("activation_batch", "activation_length", "activation_mlp")
 
+    block_size = getattr(config, "weight_block_size", None)
+
     if config.fused_mlp:
       self.wi = DenseGeneral(
           in_features_shape=in_features,
@@ -504,6 +678,7 @@ def __init__(
           mesh=self.mesh,
           use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
           debug_sharding=self.config.debug_sharding,
+          block_size=block_size,
           rngs=rngs,
       )
     else:
@@ -523,6 +698,7 @@ def __init__(
             mesh=self.mesh,
             use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
             debug_sharding=self.config.debug_sharding,
+            block_size=block_size,
             rngs=rngs,
         )
         setattr(self, dense_name, module)
@@ -541,6 +717,7 @@ def __init__(
         mesh=self.mesh,
         use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
         debug_sharding=self.config.debug_sharding,
+        block_size=block_size,
         rngs=rngs,
     )
 
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index b60ccddd08..5a03cb757c 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -516,7 +516,7 @@ def __init__(
         mesh=self.mesh,
         model_name=self.config.model_name,
         dtype=jnp.float32 if self.config.float32_gate_logits else self.dtype,
-        weight_dtype=self.weight_dtype,
+        weight_dtype=ctypes.get_weight_dtype(self.config, "gate"),
         quant=self.quant,
         kernel_init=self.kernel_init,
         kernel_axes=self.kernel_axes,
@@ -668,6 +668,66 @@ def __init__(
     else:
       self.per_expert_scale = None
 
+    if not quantizations.in_serve_mode(self.quant) and ctypes.is_fp8_dtype(self.weight_dtype):
+      scale_dtype = jnp.float32
+      block_size = getattr(self.config, "weight_block_size", None)
+      if block_size is not None:
+        if isinstance(block_size, (list, tuple)):
+          b_in = block_size[0]
+          b_out = block_size[1] if len(block_size) > 1 else block_size[0]
+        else:
+          b_in = block_size
+          b_out = block_size
+        in_blocks = self.moe_expert_input_dim // b_in if self.moe_expert_input_dim >= b_in else self.moe_expert_input_dim
+        out_blocks = moe_intermediate_dim // b_out if moe_intermediate_dim >= b_out else moe_intermediate_dim
+        if self.config.prefuse_moe_weights:
+          fused_out_dim = moe_intermediate_dim * 2
+          fused_out_blocks = fused_out_dim // b_out if fused_out_dim >= b_out else fused_out_dim
+          wi_scale_shape = (num_experts, in_blocks, fused_out_blocks)
+          wo_scale_shape = (self.num_experts, out_blocks, in_blocks)
+        else:
+          wi_scale_shape = (num_experts, in_blocks, out_blocks)
+          wo_scale_shape = (self.num_experts, out_blocks, in_blocks)
+      else:
+        wi_scale_shape = (num_experts,)
+        wo_scale_shape = (self.num_experts,)
+
+      # If using block-wise tiling, kernel scales are replicated for the block tile dimensions
+      # rather than sharded since it does not take significant memory.
+      wi_scale_sharding = (self.wi_kernel_axes[0],) + (None,) * (len(wi_scale_shape) - 1)
+      wo_scale_sharding = (self.wo_kernel_axes[0],) + (None,) * (len(wo_scale_shape) - 1)
+
+      if self.config.prefuse_moe_weights:
+        self.wi_scale = nnx.Param(
+            jnp.ones(wi_scale_shape, dtype=scale_dtype),
+            sharding=wi_scale_sharding,
+        )
+        self.wo_scale = nnx.Param(
+            jnp.ones(wo_scale_shape, dtype=scale_dtype),
+            sharding=wo_scale_sharding,
+        )
+        self.wi_0_scale = None
+        self.wi_1_scale = None
+      else:
+        self.wi_0_scale = nnx.Param(
+            jnp.ones(wi_scale_shape, dtype=scale_dtype),
+            sharding=wi_scale_sharding,
+        )
+        self.wi_1_scale = nnx.Param(
+            jnp.ones(wi_scale_shape, dtype=scale_dtype),
+            sharding=wi_scale_sharding,
+        )
+        self.wo_scale = nnx.Param(
+            jnp.ones(wo_scale_shape, dtype=scale_dtype),
+            sharding=wo_scale_sharding,
+        )
+        self.wi_scale = None
+    else:
+      self.wi_scale = None
+      self.wi_0_scale = None
+      self.wi_1_scale = None
+      self.wo_scale = None
+
     # Scale the output projection ahead of time during inference for higher generation throughput.
     if (
         self.per_expert_scale is not None
@@ -3223,21 +3283,26 @@ def __call__(
     routing_inputs = inputs if gate_inputs is None else gate_inputs.astype(gate_dtype)
     gate_logits, pre_bias_logits = self.gate(routing_inputs)
 
-    wo_kernel = jnp.asarray(self.wo[...], self.dtype)
+    wo_scale = self.wo_scale[...] if self.wo_scale is not None else None
+    wo_kernel = linears.dequantize_weight(self.wo[...], wo_scale, self.dtype)
 
     fused_kernel = None
     w0_kernel = None
     w1_kernel = None
     if cfg.prefuse_moe_weights and cfg.attention in ("vllm_rpa", "vllm_batched_rpa") and not self.is_hash_routing:
-      fused_kernel = jnp.asarray(self.wi[...], self.dtype)
+      wi_scale = self.wi_scale[...] if self.wi_scale is not None else None
+      fused_kernel = linears.dequantize_weight(self.wi[...], wi_scale, self.dtype)
     elif cfg.prefuse_moe_weights:
-      wi = jnp.asarray(self.wi[...], self.dtype)
+      wi_scale = self.wi_scale[...] if self.wi_scale is not None else None
+      wi = linears.dequantize_weight(self.wi[...], wi_scale, self.dtype)
       n = wi.shape[-1] // 2
       w0_kernel = wi[..., :n]
       w1_kernel = wi[..., n:]
     else:
-      w0_kernel = jnp.asarray(self.wi_0[...], self.dtype)
-      w1_kernel = jnp.asarray(self.wi_1[...], self.dtype)
+      wi_0_scale = self.wi_0_scale[...] if self.wi_0_scale is not None else None
+      wi_1_scale = self.wi_1_scale[...] if self.wi_1_scale is not None else None
+      w0_kernel = linears.dequantize_weight(self.wi_0[...], wi_0_scale, self.dtype)
+      w1_kernel = linears.dequantize_weight(self.wi_1[...], wi_1_scale, self.dtype)
 
     # For fused MoE path (inference only), if we have not fused expert
     # scales at init, we must apply them to wo_kernel here because
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 983c6581ba..b4ca7a1ffc 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -36,6 +36,7 @@
     MODEL_MODE_TRAIN,
     MultimodalInput,
     ShardMode,
+    get_weight_dtype,
 )
 from maxtext.layers import initializers, linears, mhc, moe, normalizations, quantizations
 from maxtext.layers import nnx_scan, nnx_wrappers
@@ -344,7 +345,7 @@ def __call__(
 
     scan_axis = self.config.param_scan_axis
     if scan_axis != 0:
-      params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params)
+      params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0) if x.ndim > scan_axis else x, params)
 
     def layer_fn(carry, scanned_vars):
       current_params, current_state = scanned_vars
@@ -368,7 +369,9 @@ def layer_fn(carry, scanned_vars):
     if scan_axis != 0:
       scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
       if scanned_params:
-        scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
+        scanned_params = jax.tree.map(
+            lambda x: jnp.moveaxis(x, 0, scan_axis) if x.ndim > scan_axis else x, scanned_params
+        )
       scanned_state = nnx.State.merge(scanned_params, scanned_other)
 
     nnx.update(self.scanned_layers, scanned_state)
@@ -420,10 +423,11 @@ def __init__(
         parameter_memory_host_offload=config.parameter_memory_host_offload,
     )
     if not config.logits_via_embedding:
+      logits_weight_dtype = get_weight_dtype(config, "logits_dense")
       self.logits_dense = linears.DenseGeneral(
           in_features_shape=config.emb_dim,
           out_features_shape=config.vocab_size,
-          weight_dtype=config.weight_dtype,
+          weight_dtype=logits_weight_dtype,
           dtype=jnp.float32 if config.logits_dot_in_fp32 else config.dtype,
           kernel_axes=("embed_vocab", "vocab"),
           shard_mode=config.shard_mode,
@@ -971,7 +975,7 @@ def _apply_layers_sequentially(
 
     scan_axis = self.config.param_scan_axis
     if scan_axis != 0:
-      params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params)
+      params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0) if x.ndim > scan_axis else x, params)
 
     layer_cls = layers.__class__
     sig = inspect.signature(layer_cls.__call__)
diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py
index 1198ef172e..4d40e2348f 100644
--- a/src/maxtext/layers/nnx_scan.py
+++ b/src/maxtext/layers/nnx_scan.py
@@ -50,7 +50,9 @@ def scan_body(carry, rng_state_slice):
   _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
 
   if param_scan_axis != 0:
-    stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), stacked_params)
+    stacked_params = jax.tree.map(
+        lambda x: jnp.moveaxis(x, 0, param_scan_axis) if x.ndim > param_scan_axis else x, stacked_params
+    )
 
   def add_scan_metadata(state, axis):
     def update_leaf(leaf):
diff --git a/src/maxtext/layers/normalizations.py b/src/maxtext/layers/normalizations.py
index 987c06fa42..0cb38012b2 100644
--- a/src/maxtext/layers/normalizations.py
+++ b/src/maxtext/layers/normalizations.py
@@ -23,7 +23,7 @@
 from jax import lax
 import jax.numpy as jnp
 from jax.sharding import NamedSharding
-from maxtext.common.common_types import Array, DType, ShardMode
+from maxtext.common.common_types import Array, DType, ShardMode, is_fp8_dtype
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import Initializer, variable_to_logically_partitioned
 from maxtext.utils import max_logging
@@ -72,7 +72,7 @@ def __init__(
     self.num_features = num_features
     self.epsilon = epsilon
     self.dtype = dtype
-    self.weight_dtype = weight_dtype
+    self.weight_dtype = dtype if is_fp8_dtype(weight_dtype) else weight_dtype
     self.shard_mode = shard_mode
     self.kernel_axes = kernel_axes
     self.scale_init = scale_init
@@ -81,7 +81,7 @@ def __init__(
     self.with_scale = with_scale
     if self.with_scale:
       self.scale = nnx.Param(
-          scale_init(rngs.params(), (num_features,), weight_dtype),
+          scale_init(rngs.params(), (num_features,), self.weight_dtype),
           out_sharding=kernel_axes,
       )
     else:
@@ -140,39 +140,80 @@ def __call__(self, x: jnp.ndarray, out_sharding: NamedSharding | None = None) ->
     return y_flat.reshape(input_shape)
 
 
-def Qwen3NextRMSNorm(
-    num_features: int,
-    epsilon: float = 1e-6,
-    dtype: DType = None,
-    weight_dtype: DType = None,
-    shard_mode=None,
-    kernel_axes=None,
-    parameter_memory_host_offload=None,
-    *,
-    rngs: nnx.Rngs,
-):
+class Qwen3NextRMSNorm(RMSNorm):
+  """RMS normalization for Qwen3 and Qwen3.5 models.
+
+  This normalization layer is specific to Qwen3 and Qwen3.5. Key characteristics:
+  1. The learnable scale parameter `scale` is initialized to ZEROS.
+  2. The scale is applied as `(1.0 + self.scale)`, making the initial scale effectively 1.0.
+  3. The normalization and scale multiplication are computed in float32 before casting
+     to self.dtype, matching Hugging Face's Qwen3 / Qwen3.5 implementation:
+     `output = (norm(x.float()) * (1.0 + weight.float())).type_as(x)`.
+     This prevents catastrophic precision loss when adding small scale offsets (~1e-3)
+     to 1.0 in bfloat16, which only has 7 mantissa bits (ULP 0.0078 at 1.0).
   """
-  Used for input and post attention layernorms
-  in Qwen3NextDecoderLayer.
 
-  This normalization layer is specific to Qwen3-Next. Key characteristics:
-  1.  The learnable scale parameter `scale` is initialized to ZEROS.
-  2.  The scale is applied as `(1.0 + self.scale)`, making the initial scale effectively 1.0.
-      This matches the PyTorch implementation of Qwen3NextRMSNorm.
+  def __init__(
+      self,
+      num_features: int,
+      epsilon: float = 1e-6,
+      dtype: Any = jnp.float32,
+      weight_dtype: Any = jnp.float32,
+      shard_mode: ShardMode = ShardMode.AUTO,
+      kernel_axes: tuple[None | str, ...] = (),
+      scale_init: Initializer = linen_initializers.zeros,
+      parameter_memory_host_offload: bool = False,
+      scale_offset: float = 1.0,
+      with_scale: bool = True,
+      *,
+      rngs: nnx.Rngs,
+  ):
+    super().__init__(
+        num_features=num_features,
+        epsilon=epsilon,
+        dtype=dtype,
+        weight_dtype=weight_dtype,
+        shard_mode=shard_mode,
+        kernel_axes=kernel_axes,
+        scale_init=scale_init,
+        parameter_memory_host_offload=parameter_memory_host_offload,
+        scale_offset=scale_offset,
+        with_scale=with_scale,
+        rngs=rngs,
+    )
 
-  """
+  def __call__(self, x: jnp.ndarray, out_sharding: NamedSharding | None = None) -> jnp.ndarray:
+    """Applies layer normalization on the input with float32 scaling."""
+    x_fp32 = jnp.asarray(x, jnp.float32)
+    mean2 = jnp.mean(lax.square(x_fp32), axis=-1, keepdims=True)
+    normed = x_fp32 * lax.rsqrt(mean2 + self.epsilon)
 
-  return nnx.data(
-      RMSNorm(
-          num_features=num_features,
-          epsilon=epsilon,
-          dtype=dtype,
-          weight_dtype=weight_dtype,
-          scale_init=linen_initializers.zeros,
-          scale_offset=1.0,
-          rngs=rngs,
-      )
-  )
+    # out_sharding must be None in auto shard mode
+    if self.shard_mode != ShardMode.EXPLICIT:
+      out_sharding = None
+
+    if out_sharding is not None:
+      out_sharding = truncate_out_sharding(out_sharding, x.ndim)
+
+    if not self.with_scale:
+      y = jnp.asarray(normed, self.dtype)
+      if out_sharding is not None:
+        y = jax.lax.with_sharding_constraint(y, out_sharding)
+      return y
+
+    scale = self.scale.get_value()
+    # Move scale to device if parameter offloading is enabled
+    if self.parameter_memory_host_offload:
+      max_logging.log("normalizations.py: Moving scale parameter to device")
+      scale = jax.device_put(scale, max_utils.device_space())
+
+    scale_fp32 = jnp.asarray(scale, jnp.float32)
+    effective_scale = scale_fp32 + self.scale_offset
+    if self.shard_mode == ShardMode.EXPLICIT:
+      effective_scale = _align_scale_with_normalized_axis(effective_scale, normed)
+
+    y = jnp.einsum("...k,k->...k", normed, effective_scale, out_sharding=out_sharding)
+    return jnp.asarray(y, self.dtype)
 
 
 class Qwen3NextRMSNormGated(nnx.Module):
@@ -199,7 +240,7 @@ def __init__(self, num_features: int, epsilon: float, dtype: DType, weight_dtype
         RMSNorm(
             num_features=num_features,
             epsilon=self.epsilon,
-            dtype=dtype,
+            dtype=jnp.float32,
             weight_dtype=weight_dtype,
             scale_init=nnx.initializers.ones,
             rngs=rngs,
@@ -273,7 +314,7 @@ def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array:
 
 
 Qwen3NextRMSNormLinen = nnx_wrappers.to_linen_class(
-    RMSNorm,
+    Qwen3NextRMSNorm,
     base_metadata_fn=variable_to_logically_partitioned,
     scale_init=linen_initializers.zeros,
     scale_offset=1.0,
diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py
index 714ec8f352..b9cf60ce6f 100644
--- a/src/maxtext/models/qwen3.py
+++ b/src/maxtext/models/qwen3.py
@@ -31,8 +31,20 @@
 from flax import linen as nn
 from flax import nnx
 
-from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE
-from maxtext.common.common_types import KV_BATCH, KV_HEAD
+from maxtext.common.common_types import (
+    AttentionType,
+    Config,
+    DType,
+    Array,
+    BATCH,
+    EMBED,
+    MODEL_MODE_TRAIN,
+    LENGTH,
+    MODEL_MODE_AUTOREGRESSIVE,
+    KV_BATCH,
+    KV_HEAD,
+    get_weight_dtype,
+)
 from maxtext.utils.sharding import (
     create_sharding,
     get_logical_axis_rules,
@@ -512,6 +524,8 @@ def __init__(
     else:
       self.cache = None  # No cache for train mode or when inputs_shape not provided
 
+    block_size = getattr(cfg, "weight_block_size", None)
+
     # Submodule instantiations
     self.in_proj_qkvz = DenseGeneral(
         in_features_shape=in_features,
@@ -520,18 +534,23 @@ def __init__(
         weight_dtype=cfg.weight_dtype,
         kernel_axes=("embed_attn", "gdn_head"),
         matmul_precision=cfg.matmul_precision,
+        block_size=block_size,
         rngs=rngs,
     )
     self.in_proj_ba = DenseGeneral(
         in_features_shape=in_features,
         out_features_shape=(self.num_v_heads * 2),
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "in_proj_ba"),
         kernel_axes=("embed_attn", "gdn_head"),
         matmul_precision=cfg.matmul_precision,
         rngs=rngs,
     )
 
+    def conv_kernel_init(key, shape, dtype=jnp.float32):
+      sample_dtype = jnp.float32 if jnp.dtype(dtype).itemsize < 2 else dtype
+      return jax.nn.initializers.lecun_normal()(key, shape, sample_dtype).astype(dtype)
+
     self.conv1d = nnx.Conv(
         in_features=conv_dim,
         out_features=conv_dim,
@@ -540,7 +559,8 @@ def __init__(
         padding="CAUSAL",
         use_bias=False,
         dtype=cfg.dtype,
-        param_dtype=cfg.weight_dtype,
+        param_dtype=get_weight_dtype(cfg, "conv1d"),
+        kernel_init=conv_kernel_init,
         precision=cfg.matmul_precision,
         rngs=rngs,
     )
@@ -548,17 +568,19 @@ def __init__(
     # Initialize A_log to match torch.log(torch.uniform(0, 16))
     def a_log_init(key, shape, dtype=jnp.float32):
       # Sample from Uniform(epsilon, 16) to avoid log(0)
-      a_vals = jax.random.uniform(key, shape=shape, dtype=dtype, minval=1e-9, maxval=16.0)
-      return jnp.log(a_vals)
+      a_vals = jax.random.uniform(key, shape=shape, dtype=jnp.float32, minval=1e-9, maxval=16.0)
+      return jnp.log(a_vals).astype(dtype)
 
-    self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype))
-    self.dt_bias = nnx.Param(nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype))
+    self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=get_weight_dtype(cfg, "A_log")))
+    self.dt_bias = nnx.Param(
+        nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=get_weight_dtype(cfg, "dt_bias"))
+    )
 
     self.norm = Qwen3NextRMSNormGated(
         num_features=self.head_v_dim,  # Normalize over the head dimension (D_v)
         epsilon=cfg.normalization_layer_epsilon,
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "norm"),
         rngs=rngs,
     )
     self.out_proj = DenseGeneral(
@@ -568,6 +590,7 @@ def a_log_init(key, shape, dtype=jnp.float32):
         weight_dtype=cfg.weight_dtype,
         kernel_axes=("gdn_head", "embed_attn"),
         matmul_precision=cfg.matmul_precision,
+        block_size=block_size,
         rngs=rngs,
     )
 
@@ -1165,6 +1188,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn
           out_features_shape=1,
           use_bias=False,  # Qwen3-Next shared_expert_gate does not have a bias
           dtype=cfg.dtype,
+          weight_dtype=get_weight_dtype(cfg, "shared_expert_gate"),
           kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"),
           kernel_axes=("embed", None),
           matmul_precision=cfg.matmul_precision,
@@ -1322,7 +1346,7 @@ def __init__(
         num_features=cfg.emb_dim,
         epsilon=cfg.normalization_layer_epsilon,
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "norm"),
         rngs=rngs,
     )
 
@@ -1351,7 +1375,7 @@ def __init__(
         num_features=cfg.emb_dim,
         epsilon=cfg.normalization_layer_epsilon,
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "norm"),
         rngs=rngs,
     )
 
diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py
index 331df17f41..edf6233eea 100644
--- a/src/maxtext/models/qwen3_5.py
+++ b/src/maxtext/models/qwen3_5.py
@@ -24,7 +24,7 @@
 from flax import linen as nn
 from flax import nnx
 
-from maxtext.common.common_types import Config, Array
+from maxtext.common.common_types import Config, Array, get_weight_dtype
 from maxtext.layers import initializers as max_initializers
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.normalizations import Qwen3NextRMSNorm
@@ -138,7 +138,7 @@ def __init__(
         num_features=cfg.emb_dim,
         epsilon=cfg.normalization_layer_epsilon,
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "norm"),
         rngs=rngs,
     )
 
@@ -167,7 +167,7 @@ def __init__(
         num_features=cfg.emb_dim,
         epsilon=cfg.normalization_layer_epsilon,
         dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
+        weight_dtype=get_weight_dtype(cfg, "norm"),
         rngs=rngs,
     )
 
diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py
index cf5da0edcb..061c237239 100644
--- a/src/maxtext/multimodal/processor.py
+++ b/src/maxtext/multimodal/processor.py
@@ -36,6 +36,8 @@
     "qwen3-vl-4b": ("qwen3_vl", "qwen3"),
     "qwen3-vl-30b-a3b": ("qwen3_vl", "qwen3_moe"),
     "qwen3.5-35b-a3b": ("qwen3_5", "qwen3_5"),
+    "qwen3.5-35b-a3b-fp8": ("qwen3_5", "qwen3_5"),
+    "qwen3.5-35b-fp8": ("qwen3_5", "qwen3_5"),
     "qwen3.5-397b-a17b": ("qwen3_5", "qwen3_5"),
     # Stitched model
     "maxtext-omni-gemma3-qwen3": ("gemma3", "qwen3"),
diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py
index 30f6e65124..5e34db776d 100644
--- a/src/maxtext/utils/globals.py
+++ b/src/maxtext/utils/globals.py
@@ -84,7 +84,10 @@
     "qwen3-omni-30b-a3b": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
     "qwen3-next-80b-a3b": "Qwen/Qwen3-Next-80B-A3B-Instruct",
     "qwen3.5-397b-a17b": "Qwen/Qwen3.5-397B-A17B",
+    "qwen3.5-397b-a17b-fp8": "Qwen/Qwen3.5-397B-A17B-FP8",
     "qwen3.5-35b-a3b": "Qwen/Qwen3.5-35B-A3B",
+    "qwen3.5-35b-a3b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8",
+    "qwen3.5-35b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8",
     "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "mistral-7b": "mistralai/Mistral-7B-v0.1",
     "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1",
diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py
index 1b67c2592c..261b00dcc0 100644
--- a/src/maxtext/utils/maxtext_utils_nnx.py
+++ b/src/maxtext/utils/maxtext_utils_nnx.py
@@ -218,6 +218,7 @@ def nnx_update_sharding_meta(variable, transform_fn):
   return variable
 
 
+
 def nnx_remove_scan_axis(tree, name="layers"):
   """Removes the given scan axis from the PartitionSpec."""
 
diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py
index 9440d64821..9dce2626ba 100644
--- a/tests/unit/checkpointing_test.py
+++ b/tests/unit/checkpointing_test.py
@@ -17,15 +17,19 @@
 import asyncio
 import json
 import os
+import tempfile
 from unittest import mock
 
 from absl.testing import absltest
 from absl.testing import parameterized
 from etils import epath
+from flax import nnx
 from flax.training import train_state
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from maxtext.layers import linears
+import orbax.checkpoint as ocp
 from maxtext.checkpoint_conversion.utils import load_dynamic
 from maxtext.checkpoint_conversion.utils.tensor_handling import (
     _binary_chunked_stack,
@@ -512,5 +516,243 @@ def test_error_handler_raises_runtime_error(self):
       self.assertIs(cm.exception.__cause__, original_error)
 
 
+class FP8DequantizeOnLoadTest(parameterized.TestCase):
+  """Tests for dequantize-on-load parameter restoration."""
+
+  def setUp(self):
+    super().setUp()
+    self.tmp_dir = tempfile.TemporaryDirectory()
+
+  def tearDown(self):
+    self.tmp_dir.cleanup()
+    super().tearDown()
+
+  def test_load_fp8_checkpoint_into_bf16_nnx_model(self):
+    """Loading an FP8 checkpoint into a BF16 NNX model restores dequantized BF16 weights and drops scale."""
+    class BF16Model(nnx.Module):
+
+      def __init__(self, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(4, 2, rngs=rngs, dtype=jnp.bfloat16, param_dtype=jnp.bfloat16)
+
+    model = BF16Model(rngs=nnx.Rngs(0))
+    _, params_abstract, _ = nnx.split(model, nnx.Param, ...)
+
+    fp8_kernel = jnp.array([[0.25, 0.5], [1.0, 1.5], [0.125, 0.75], [2.0, 0.5]], dtype=jnp.float8_e4m3fn)
+    scale = jnp.array(2.0, dtype=jnp.float32)
+    bias = jnp.zeros((2,), dtype=jnp.bfloat16)
+
+    ckpt_weights = {
+        "linear": {
+            "kernel": fp8_kernel,
+            "kernel_scale": scale,
+            "bias": bias,
+        }
+    }
+
+    path = os.path.join(self.tmp_dir.name, "fp8_ckpt")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {"params": {"params": ckpt_weights}},
+        force=True,
+    )
+
+    expected_kernel = linears.dequantize_weight(fp8_kernel, scale, compute_dtype=jnp.bfloat16)
+
+    restored = checkpointing.load_params_from_path(path, params_abstract, 8)
+    self.assertIsInstance(restored, nnx.State)
+    pure = restored.to_pure_dict()
+
+    self.assertNotIn("kernel_scale", pure["linear"])
+    self.assertEqual(pure["linear"]["kernel"].dtype, jnp.bfloat16)
+    self.assertEqual(pure["linear"]["kernel"].shape, (4, 2))
+    np.testing.assert_allclose(
+        np.array(pure["linear"]["kernel"]),
+        np.array(expected_kernel),
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+  def test_load_fp8_checkpoint_into_fp8_nnx_model(self):
+    """Loading an FP8 checkpoint into an FP8 NNX model preserves FP8 kernel and scale."""
+    class FP8Model(nnx.Module):
+
+      def __init__(self, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(4, 2, rngs=rngs, dtype=jnp.bfloat16, param_dtype=jnp.float8_e4m3fn)
+        self.linear.kernel_scale = nnx.Param(jnp.ones((), dtype=jnp.float32))
+
+    model = FP8Model(rngs=nnx.Rngs(0))
+    _, params_abstract, _ = nnx.split(model, nnx.Param, ...)
+
+    fp8_kernel = jnp.array([[0.25, 0.5], [1.0, 1.5], [0.125, 0.75], [2.0, 0.5]], dtype=jnp.float8_e4m3fn)
+    scale = jnp.array(3.5, dtype=jnp.float32)
+    bias = jnp.zeros((2,), dtype=jnp.float8_e4m3fn)
+
+    ckpt_weights = {
+        "linear": {
+            "kernel": fp8_kernel,
+            "kernel_scale": scale,
+            "bias": bias,
+        }
+    }
+
+    path = os.path.join(self.tmp_dir.name, "fp8_to_fp8_ckpt")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {"params": {"params": ckpt_weights}},
+        force=True,
+    )
+
+    restored = checkpointing.load_params_from_path(path, params_abstract, 8)
+    self.assertIsInstance(restored, nnx.State)
+    pure = restored.to_pure_dict()
+
+    self.assertIn("kernel_scale", pure["linear"])
+    self.assertEqual(pure["linear"]["kernel"].dtype, jnp.float8_e4m3fn)
+    self.assertEqual(pure["linear"]["kernel_scale"].dtype, jnp.float32)
+    np.testing.assert_array_equal(np.array(pure["linear"]["kernel"]), np.array(fp8_kernel))
+    np.testing.assert_array_equal(np.array(pure["linear"]["kernel_scale"]), np.array(scale))
+
+  def test_load_fp8_checkpoint_into_bf16_linen_dict(self):
+    """Loading an FP8 checkpoint into a BF16 Linen parameter dict restores dequantized BF16 weights."""
+    target_weights = {
+        "params": {
+            "linear": {
+                "kernel": jax.ShapeDtypeStruct(shape=(4, 2), dtype=jnp.bfloat16),
+                "bias": jax.ShapeDtypeStruct(shape=(2,), dtype=jnp.bfloat16),
+            }
+        }
+    }
+
+    fp8_kernel = jnp.array([[0.5, 1.0], [0.25, 0.75], [1.5, 0.125], [0.5, 2.0]], dtype=jnp.float8_e4m3fn)
+    scale = jnp.array(0.5, dtype=jnp.float32)
+    bias = jnp.zeros((2,), dtype=jnp.bfloat16)
+
+    ckpt_weights = {
+        "params": {
+            "linear": {
+                "kernel": fp8_kernel,
+                "kernel_scale": scale,
+                "bias": bias,
+            }
+        }
+    }
+
+    path = os.path.join(self.tmp_dir.name, "fp8_linen_ckpt")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {"params": ckpt_weights},
+        force=True,
+    )
+
+    expected_kernel = linears.dequantize_weight(fp8_kernel, scale, compute_dtype=jnp.bfloat16)
+
+    restored = checkpointing.load_params_from_path(path, target_weights, 8)
+    self.assertNotIsInstance(restored, nnx.State)
+    self.assertIn("params", restored)
+    self.assertNotIn("kernel_scale", restored["params"]["linear"])
+    self.assertEqual(restored["params"]["linear"]["kernel"].dtype, jnp.bfloat16)
+    np.testing.assert_allclose(
+        np.array(restored["params"]["linear"]["kernel"]),
+        np.array(expected_kernel),
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+  def test_load_fp8_checkpoint_with_per_channel_scale_into_bf16_model(self):
+    """Loading an FP8 checkpoint with per-channel scale dequantizes properly."""
+    class BF16Model(nnx.Module):
+
+      def __init__(self, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(4, 2, rngs=rngs, dtype=jnp.bfloat16, param_dtype=jnp.bfloat16)
+
+    model = BF16Model(rngs=nnx.Rngs(0))
+    _, params_abstract, _ = nnx.split(model, nnx.Param, ...)
+
+    fp8_kernel = jnp.array([[0.25, 0.5], [1.0, 1.5], [0.125, 0.75], [2.0, 0.5]], dtype=jnp.float8_e4m3fn)
+    scale = jnp.array([2.0, 4.0], dtype=jnp.float32)
+    bias = jnp.zeros((2,), dtype=jnp.bfloat16)
+
+    ckpt_weights = {
+        "linear": {
+            "kernel": fp8_kernel,
+            "kernel_scale": scale,
+            "bias": bias,
+        }
+    }
+
+    path = os.path.join(self.tmp_dir.name, "fp8_channel_scale_ckpt")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {"params": {"params": ckpt_weights}},
+        force=True,
+    )
+
+    expected_kernel = linears.dequantize_weight(fp8_kernel, scale, compute_dtype=jnp.bfloat16)
+
+    restored = checkpointing.load_params_from_path(path, params_abstract, 8)
+    pure = restored.to_pure_dict()
+
+    self.assertNotIn("kernel_scale", pure["linear"])
+    self.assertEqual(pure["linear"]["kernel"].dtype, jnp.bfloat16)
+    np.testing.assert_allclose(
+        np.array(pure["linear"]["kernel"]),
+        np.array(expected_kernel),
+        rtol=1e-3,
+        atol=1e-3,
+    )
+
+  def test_load_fp8_checkpoint_with_moe_scales_into_bf16_model(self):
+    """Loading an FP8 checkpoint with MoE expert weights (wi_0, wi_1, wo) dequantizes properly."""
+    class BF16MoEModel(nnx.Module):
+
+      def __init__(self, rngs: nnx.Rngs):
+        self.wi_0 = nnx.Param(jnp.zeros((2, 4, 8), dtype=jnp.bfloat16))
+        self.wi_1 = nnx.Param(jnp.zeros((2, 4, 8), dtype=jnp.bfloat16))
+        self.wo = nnx.Param(jnp.zeros((2, 8, 4), dtype=jnp.bfloat16))
+
+    model = BF16MoEModel(rngs=nnx.Rngs(0))
+    _, params_abstract, _ = nnx.split(model, nnx.Param, ...)
+
+    fp8_wi_0 = jnp.array(np.random.randn(2, 4, 8), dtype=jnp.float8_e4m3fn)
+    wi_0_scale = jnp.array(np.random.rand(2, 1, 1), dtype=jnp.float32)
+    fp8_wi_1 = jnp.array(np.random.randn(2, 4, 8), dtype=jnp.float8_e4m3fn)
+    wi_1_scale = jnp.array(np.random.rand(2, 1, 1), dtype=jnp.float32)
+    fp8_wo = jnp.array(np.random.randn(2, 8, 4), dtype=jnp.float8_e4m3fn)
+    wo_scale = jnp.array(np.random.rand(2, 1, 1), dtype=jnp.float32)
+
+    ckpt_weights = {
+        "wi_0": fp8_wi_0,
+        "wi_0_scale": wi_0_scale,
+        "wi_1": fp8_wi_1,
+        "wi_1_scale": wi_1_scale,
+        "wo": fp8_wo,
+        "wo_scale": wo_scale,
+    }
+
+    path = os.path.join(self.tmp_dir.name, "fp8_moe_scale_ckpt")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {"params": {"params": ckpt_weights}},
+        force=True,
+    )
+
+    expected_wi_0 = linears.dequantize_weight(fp8_wi_0, wi_0_scale, compute_dtype=jnp.bfloat16)
+    expected_wi_1 = linears.dequantize_weight(fp8_wi_1, wi_1_scale, compute_dtype=jnp.bfloat16)
+    expected_wo = linears.dequantize_weight(fp8_wo, wo_scale, compute_dtype=jnp.bfloat16)
+
+    restored = checkpointing.load_params_from_path(path, params_abstract, 8)
+    pure = restored.to_pure_dict()
+
+    self.assertNotIn("wi_0_scale", pure)
+    self.assertNotIn("wi_1_scale", pure)
+    self.assertNotIn("wo_scale", pure)
+    self.assertEqual(pure["wi_0"].dtype, jnp.bfloat16)
+    self.assertEqual(pure["wi_1"].dtype, jnp.bfloat16)
+    self.assertEqual(pure["wo"].dtype, jnp.bfloat16)
+    np.testing.assert_allclose(np.array(pure["wi_0"]), np.array(expected_wi_0), rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(np.array(pure["wi_1"]), np.array(expected_wi_1), rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(np.array(pure["wo"]), np.array(expected_wo), rtol=1e-3, atol=1e-3)
+
+
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/unit/configs_test.py b/tests/unit/configs_test.py
index 2a7bd0f660..46a56184a2 100644
--- a/tests/unit/configs_test.py
+++ b/tests/unit/configs_test.py
@@ -246,6 +246,10 @@ def test_mistral_configs(config_file):
     os.path.join(CONFIGS_DIR, "models", "qwen3-480b-a35b.yml"),
     os.path.join(CONFIGS_DIR, "models", "qwen3-next-80b-a3b.yml"),
     os.path.join(CONFIGS_DIR, "models", "qwen3-omni-30b-a3b.yml"),
+    os.path.join(CONFIGS_DIR, "models", "qwen3.5-35b-a3b.yml"),
+    os.path.join(CONFIGS_DIR, "models", "qwen3.5-35b-a3b-fp8.yml"),
+    os.path.join(CONFIGS_DIR, "models", "qwen3.5-397b-a17b.yml"),
+    os.path.join(CONFIGS_DIR, "models", "qwen3.5-397b-a17b-fp8.yml"),
 ]
 
 
diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py
index a20d29a0d9..793997c119 100644
--- a/tests/unit/linears_test.py
+++ b/tests/unit/linears_test.py
@@ -207,6 +207,164 @@ def test_axis_1(self):
   def test_axis_0(self):
     self._run_dense_test(0, 2, (3, 4, 8))
 
+  def test_dequantize_weight_scalar(self):
+    w = jnp.ones((4, 8), dtype=jnp.float8_e4m3fn)
+    scale = jnp.array(0.5, dtype=jnp.float32)
+    w_dequant = linears.dequantize_weight(w, scale, compute_dtype=jnp.bfloat16)
+    self.assertEqual(w_dequant.shape, (4, 8))
+    self.assertEqual(w_dequant.dtype, jnp.bfloat16)
+    np.testing.assert_allclose(w_dequant, np.full((4, 8), 0.5, dtype=np.float32), rtol=1e-3)
+
+  def test_dequantize_weight_per_channel(self):
+    w = jnp.ones((4, 8), dtype=jnp.float8_e4m3fn)
+    scale = jnp.arange(1, 9, dtype=jnp.float32)
+    w_dequant = linears.dequantize_weight(w, scale, compute_dtype=jnp.bfloat16)
+    self.assertEqual(w_dequant.shape, (4, 8))
+    expected = np.tile(np.arange(1, 9, dtype=np.float32), (4, 1))
+    np.testing.assert_allclose(w_dequant, expected, rtol=1e-3)
+
+  def test_dequantize_weight_block_wise(self):
+    w = jnp.ones((4, 8), dtype=jnp.float8_e4m3fn)
+    scale = jnp.array([[2.0, 3.0], [4.0, 5.0]], dtype=jnp.float32)  # 2x2 blocks of size 2x4
+    w_dequant = linears.dequantize_weight(w, scale, compute_dtype=jnp.bfloat16)
+    self.assertEqual(w_dequant.shape, (4, 8))
+    expected = np.block([[np.full((2, 4), 2.0), np.full((2, 4), 3.0)], [np.full((2, 4), 4.0), np.full((2, 4), 5.0)]])
+    np.testing.assert_allclose(w_dequant, expected, rtol=1e-3)
+
+  def test_fp8_e4m3fn_dense_general(self):
+    batch_size = 2
+    in_features = 4
+    out_features = 8
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=out_features,
+        weight_dtype=jnp.float8_e4m3fn,
+        dtype=jnp.bfloat16,
+        rngs=self.rngs,
+    )
+
+    self.assertEqual(layer.kernel[...].dtype, jnp.float8_e4m3fn)
+    self.assertIsNotNone(layer.kernel_scale)
+    self.assertEqual(layer.kernel_scale[...].shape, ())
+
+    inputs = jnp.ones((batch_size, in_features), dtype=jnp.bfloat16)
+    outputs = layer(inputs)
+
+    self.assertEqual(outputs.shape, (batch_size, out_features))
+    self.assertEqual(outputs.dtype, jnp.bfloat16)
+
+  def test_fp8_e5m2_dense_general(self):
+    batch_size = 2
+    in_features = 4
+    out_features = 8
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=out_features,
+        weight_dtype=jnp.float8_e5m2,
+        dtype=jnp.bfloat16,
+        rngs=self.rngs,
+    )
+
+    self.assertEqual(layer.kernel[...].dtype, jnp.float8_e5m2)
+    self.assertIsNotNone(layer.kernel_scale)
+
+    inputs = jnp.ones((batch_size, in_features), dtype=jnp.bfloat16)
+    outputs = layer(inputs)
+
+    self.assertEqual(outputs.shape, (batch_size, out_features))
+    self.assertEqual(outputs.dtype, jnp.bfloat16)
+
+  def test_fp8_block_wise_scale(self):
+    in_features = 128
+    out_features = 256
+
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=out_features,
+        weight_dtype=jnp.float8_e4m3fn,
+        dtype=jnp.bfloat16,
+        block_size=64,
+        kernel_axes=("embed", "mlp"),
+        rngs=self.rngs,
+    )
+
+    self.assertEqual(layer.kernel[...].shape, (128, 256))
+    self.assertEqual(layer.kernel_scale[...].shape, (2, 4))
+    # Block-compressed scale dims (2, 4) are far smaller than the kernel's (128, 256), so
+    # they're replicated rather than inheriting the kernel's mesh axes.
+    self.assertEqual(layer.scale_axes, (None, None))
+
+    inputs = jnp.ones((2, in_features), dtype=jnp.bfloat16)
+    outputs = layer(inputs)
+    self.assertEqual(outputs.shape, (2, out_features))
+
+  def test_kernel_scale_sharding_inference(self):
+    # Test block scale sharding
+    layer_block = linears.DenseGeneral(
+        in_features_shape=64,
+        out_features_shape=128,
+        weight_dtype=jnp.float8_e4m3fn,
+        block_size=32,
+        kernel_axes=("embed", "mlp"),
+        rngs=self.rngs,
+    )
+    # Block-compressed (2, 4) is replicated, not sharded on the kernel's ("embed", "mlp") axes.
+    self.assertEqual(layer_block.scale_axes, (None, None))
+
+    # Test scalar scale sharding
+    layer_scalar = linears.DenseGeneral(
+        in_features_shape=64,
+        out_features_shape=128,
+        weight_dtype=jnp.float8_e4m3fn,
+        kernel_axes=("embed", "mlp"),
+        rngs=self.rngs,
+    )
+    self.assertEqual(layer_scalar.scale_axes, ())
+
+    # Test per-channel scale sharding
+    layer_channel = linears.DenseGeneral(
+        in_features_shape=64,
+        out_features_shape=128,
+        weight_dtype=jnp.float8_e4m3fn,
+        scale_shape=(1, 128),
+        kernel_axes=("embed", "mlp"),
+        rngs=self.rngs,
+    )
+    self.assertEqual(layer_channel.scale_axes, (None, "mlp"))
+
+  def test_fp8_slice_bounds(self):
+    in_features = 64
+    out_features = 128
+    layer = linears.DenseGeneral(
+        in_features_shape=in_features,
+        out_features_shape=out_features,
+        weight_dtype=jnp.float8_e4m3fn,
+        dtype=jnp.bfloat16,
+        block_size=32,
+        rngs=self.rngs,
+    )
+    inputs = jnp.ones((2, in_features), dtype=jnp.bfloat16)
+    sliced_outputs = layer(inputs, slice_bounds=(0, 32))
+    self.assertEqual(sliced_outputs.shape, (2, 32))
+    self.assertEqual(sliced_outputs.dtype, jnp.bfloat16)
+
+    full_outputs = layer(inputs)
+    np.testing.assert_allclose(sliced_outputs, full_outputs[:, :32], rtol=1e-3, atol=1e-3)
+
+  def test_dequantize_weight_block_divisibility(self):
+    w = jnp.ones((64, 128), dtype=jnp.float8_e4m3fn)
+    scale_valid = jnp.ones((2, 4), dtype=jnp.float32) * 2.0
+    dequant = linears.dequantize_weight(w, scale_valid, compute_dtype=jnp.bfloat16)
+    self.assertEqual(dequant.shape, (64, 128))
+    self.assertEqual(dequant.dtype, jnp.bfloat16)
+    np.testing.assert_allclose(dequant, 2.0, rtol=1e-3, atol=1e-3)
+
+    scale_invalid = jnp.ones((3, 5), dtype=jnp.float32)
+    with self.assertRaises(ValueError):
+      linears.dequantize_weight(w, scale_invalid, compute_dtype=jnp.bfloat16)
+
 
 class MlpBlockTest(unittest.TestCase):
   """Tests for MlpBlock."""
diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py
index c2832f4ce9..ca3b49deed 100644
--- a/tests/unit/nnx_decoders_test.py
+++ b/tests/unit/nnx_decoders_test.py
@@ -1318,5 +1318,64 @@ def __call__(self, x, **kwargs):
       maxtext_utils_nnx.nnx_add_and_sync_scan_axis = original_add_scan_axis
 
 
+class TestNNXDecoderFP8WeightOnly(unittest.TestCase):
+  """Tests for NNXDecoder with FP8 weight-only storage and dynamic dequantization."""
+
+  def setUp(self):
+    super().setUp()
+    self.cfg = _make_config(
+        weight_dtype="float8_e4m3fn",
+        dtype="bfloat16",
+        unquantized_modules=["token_embedder", "logits_dense"],
+    )
+    self.mesh = _make_mesh(self.cfg)
+    self.rngs = nnx.Rngs(params=0, dropout=1)
+    self.decoder = NNXDecoder(
+        config=self.cfg,
+        mesh=self.mesh,
+        rngs=self.rngs,
+    )
+    self.shared_embedding = Embed(
+        num_embeddings=self.cfg.vocab_size,
+        num_features=self.cfg.emb_dim,
+        dtype=self.cfg.dtype,
+        embedding_init=jax.nn.initializers.normal(stddev=1.0),
+        config=self.cfg,
+        mesh=self.mesh,
+        rngs=self.rngs,
+    )
+
+  def test_fp8_weights_and_unquantized_layers(self):
+    """Verifies that dense linear weights are FP8 while embedding and norms are BF16."""
+    layer_0 = self.decoder.layers_0
+    self.assertEqual(layer_0.self_attention.query.kernel[...].dtype, jnp.float8_e4m3fn)
+    self.assertIsNotNone(layer_0.self_attention.query.kernel_scale)
+    self.assertEqual(layer_0.mlp.wi_0.kernel[...].dtype, jnp.float8_e4m3fn)
+    self.assertEqual(self.shared_embedding.embedding[...].dtype, jnp.bfloat16)
+    self.assertEqual(self.decoder.decoder_norm.scale[...].dtype, jnp.bfloat16)
+
+  def test_fp8_forward_pass_execution(self):
+    """Verifies that an end-to-end forward pass with dynamic FP8 dequantization executes and produces valid logits."""
+    cfg = self.cfg
+    batch = cfg.global_batch_size_to_train_on
+    seq_len = cfg.max_target_length
+    ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq_len), 0, cfg.vocab_size)
+    segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR)
+    positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len))
+
+    logits, hidden_state, _ = self.decoder(
+        self.shared_embedding,
+        ids,
+        positions,
+        decoder_segment_ids=segment_ids,
+        deterministic=True,
+        model_mode=MODEL_MODE_TRAIN,
+    )
+
+    self.assertEqual(logits.shape, (batch, seq_len, cfg.vocab_size))
+    self.assertEqual(hidden_state.shape, (batch, seq_len, cfg.emb_dim))
+    self.assertTrue(jnp.all(jnp.isfinite(logits)))
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py
index a51b23980f..392c71d355 100644
--- a/tests/utils/forward_pass_logit_checker.py
+++ b/tests/utils/forward_pass_logit_checker.py
@@ -73,6 +73,8 @@
 import argparse
 import functools
 import os
+
+os.environ.setdefault("HF_HOME", "/dev/shm/hf_cache")
 from pathlib import Path
 import sys
 import absl
diff --git a/tests/utils/run_fp8_logit_test.py b/tests/utils/run_fp8_logit_test.py
new file mode 100644
index 0000000000..1013063201
--- /dev/null
+++ b/tests/utils/run_fp8_logit_test.py
@@ -0,0 +1,29 @@
+import os
+import sys
+import runpy
+
+sys.path.insert(0, os.getcwd())
+os.environ.setdefault("HF_HOME", "/dev/shm/hf_cache")
+
+sys.argv = [
+    "forward_pass_logit_checker.py",
+    "src/maxtext/configs/base.yml",
+    "model_name=qwen3.5-35b-a3b-fp8",
+    "load_parameters_path=/dev/shm/maxtext_qwen3.5_35b_fp8/0/items",
+    "scan_layers=true",
+    "per_device_batch_size=1",
+    "max_prefill_predict_length=4",
+    "max_target_length=4",
+    "async_checkpointing=false",
+    "sparse_matmul=false",
+    "ici_fsdp_parallelism=1",
+    "ici_expert_parallelism=-1",
+    "matmul_precision=highest",
+    "float32_logits=true",
+    "float32_qk_product=true",
+    "--golden_logits_path=/dev/shm/golden_qwen3.5_35b_fp8.pkl",
+    "--max_kl_div=0.2",
+]
+
+print("Launching forward pass logit test against /dev/shm/golden_qwen3.5_35b_fp8.pkl with --max_kl_div=0.2...")
+runpy.run_module("tests.utils.forward_pass_logit_checker", run_name="__main__")