[Fix] Support Kimi-K3 ModelOpt mixed NVFP4/FP8 checkpoint (#35077)

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -517,6 +517,7 @@ def _compute_g1_scale_c(
     g1_alphas: torch.Tensor,
     g1_alphas_up: torch.Tensor,
     is_gated: bool,
+    activation: Optional[str] = None,
 ) -> torch.Tensor:
     """TRT-LLM GEMM1-output scale for the up (w3) half.
 
@@ -526,6 +527,11 @@ def _compute_g1_scale_c(
     scale passes g1_alphas as g1_alphas_up and recovers the single-scale value;
     non-gated (Relu2) has no gate half, so it is just 1/a2_scale per expert.
     """
+    if activation == "situ":
+        # SiTU consumes both GEMM1 scales before tanh; scale_c carries only
+        # the GEMM2 input requantization factor.
+        num_experts = g1_alphas.shape[0]
+        return w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous()
     if is_gated:
         return (w2_input_scale_quant * g1_alphas_up).to(torch.float32)
     num_experts = g1_alphas.shape[0]
@@ -596,7 +602,11 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
     g1_alphas = cast(torch.Tensor, layer.g1_alphas)
     g1_alphas_up = cast(torch.Tensor, getattr(layer, "g1_alphas_up", g1_alphas))
     g1_scale_c = _compute_g1_scale_c(
-        w2_input_scale_quant, g1_alphas, g1_alphas_up, layer.moe_runner_config.is_gated
+        w2_input_scale_quant,
+        g1_alphas,
+        g1_alphas_up,
+        layer.moe_runner_config.is_gated,
+        activation=layer.moe_runner_config.activation,
     )
     copy_or_rebind_param(layer, "g1_scale_c", g1_scale_c)
 
@@ -612,6 +622,7 @@ def get_activation_type(activation: str, is_gated: bool = True) -> int:
         _ACTIVATION_STR_TO_TYPE = {
             "silu": ActivationType.Swiglu,
             "gelu": ActivationType.Geglu,
+            "situ": ActivationType.Situ,
         }
     else:
         _ACTIVATION_STR_TO_TYPE = {
@@ -956,7 +967,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
     from sglang.srt.layers.moe.topk import TopKOutputChecker
     from sglang.srt.layers.moe.utils import RoutingMethodType
 
-    _SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu"}
+    _SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu", "situ"}
     assert runner_config.activation in _SUPPORTED_FP4_ACTIVATIONS, (
         f"Only {_SUPPORTED_FP4_ACTIVATIONS} are supported for FP4 MoE, "
         f"got '{runner_config.activation}'."

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -776,6 +776,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
         packed_modules_mapping: Optional[Dict[str, List[str]]],
         quantized_layers: Dict[str, Dict[str, Any]],
         fp8_config: ModelOptFp8Config,
+        fp8_pb_wo_config: Fp8Config,
         nvfp4_config: ModelOptFp4Config,
         nvfp4a16_config: ModelOptFp4Config,
         mxfp8_config: Fp8Config,
@@ -783,6 +784,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
         super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
         self.quantized_layers = quantized_layers
         self.fp8_config = fp8_config
+        self.fp8_pb_wo_config = fp8_pb_wo_config
         self.mxfp8_config = mxfp8_config
         self.nvfp4_config = nvfp4_config
         self.nvfp4a16_config = nvfp4a16_config
@@ -867,6 +869,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             exclude_modules=[],
             packed_modules_mapping=packed_modules_mapping,
         )
+        fp8_pb_wo_config = Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            activation_scheme="dynamic",
+            weight_block_size=[128, 128],
+            packed_modules_mapping=packed_modules_mapping,
+        )
         mxfp8_config = Fp8Config(
             is_checkpoint_fp8_serialized=True,
             activation_scheme="dynamic",
@@ -896,6 +904,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             packed_modules_mapping=packed_modules_mapping,
             quantized_layers=quantized_layers,
             fp8_config=fp8_config,
+            fp8_pb_wo_config=fp8_pb_wo_config,
             mxfp8_config=mxfp8_config,
             nvfp4_config=nvfp4_config,
             nvfp4a16_config=nvfp4a16_config,
@@ -977,6 +986,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
                 return UnquantizedLinearMethod()
             if quant_algo == "FP8":
                 return ModelOptFp8LinearMethod(self.fp8_config)
+            if quant_algo == "FP8_PB_WO":
+                return Fp8LinearMethod(self.fp8_pb_wo_config)
             if quant_algo == "MXFP8":
                 return Fp8LinearMethod(self.mxfp8_config)
             if quant_algo == "NVFP4":
@@ -2511,9 +2522,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         )
 
         if layer.moe_runner_config.is_gated and self.enable_flashinfer_trtllm_moe:
+            runner_config = layer.moe_runner_config
+            is_situ = runner_config.activation == "situ"
             gemm1_clamp_limit = (
-                layer.moe_runner_config.gemm1_clamp_limit
+                None
-                or layer.moe_runner_config.swiglu_limit
+                if is_situ
+                else (runner_config.gemm1_clamp_limit or runner_config.swiglu_limit)
             )
             if gemm1_clamp_limit is not None:
                 copy_or_rebind_param(
@@ -2522,21 +2536,26 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
                     (gemm1_clamp_limit / layer.g1_alphas).to(torch.float32),
                 )
 
-            if layer.moe_runner_config.gemm1_alpha is not None:
+            if runner_config.gemm1_alpha is not None:
                 copy_or_rebind_param(
                     layer,
                     "gemm1_alpha",
                     torch.full_like(
                         layer.g1_alphas,
-                        layer.moe_runner_config.gemm1_alpha,
+                        runner_config.gemm1_alpha,
                         dtype=torch.float32,
                     ),
                 )
-                copy_or_rebind_param(
+                gemm1_beta = (
-                    layer,
+                    torch.full_like(
-                    "gemm1_beta",
+                        layer.g1_alphas,
-                    (1.0 / layer.g1_alphas).to(torch.float32),
+                        runner_config.gemm1_clamp_limit,
+                        dtype=torch.float32,
                     )
+                    if is_situ
+                    else (1.0 / layer.g1_alphas).to(torch.float32)
+                )
+                copy_or_rebind_param(layer, "gemm1_beta", gemm1_beta)
 
         # TODO: for flashinfer always do MOE_NVFP4_DISPATCH
         use_dispatch_fp4 = not self.quant_config.use_per_token_activation and (
@@ -2802,9 +2821,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
             self, "_moe_runner_backend", get_moe_runner_backend()
         )
 
-        assert (
+        assert activation in _SUPPORTED_ACT_STRS or (
-            activation in _SUPPORTED_ACT_STRS
+            activation == "situ" and moe_runner_backend.is_flashinfer_trtllm()
-        ), f"{activation=} not in supported {_SUPPORTED_ACT_STRS}"
+        ), f"{activation=} is unsupported by {moe_runner_backend}"
         moe_runner_config = self.moe_runner_config
 
         if moe_runner_backend.is_marlin():

python/sglang/srt/models/kimi_k3.py

@@ -73,6 +73,7 @@ from sglang.srt.layers.moe.utils import (
     get_moe_runner_backend,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import (
@@ -142,6 +143,35 @@ def _cdiv(a: int, b: int) -> int:
     return (a + b - 1) // b
 
 
+def _uses_modelopt_fp8_pb_wo(
+    quant_config: Optional[QuantizationConfig], prefix: str
+) -> bool:
+    resolver = getattr(quant_config, "_resolve_quant_algo", None)
+    return resolver is not None and resolver(prefix) == "FP8_PB_WO"
+
+
+def _maybe_map_fp8_pb_scale_name(name: str, params_dict: dict) -> str:
+    """Map ModelOpt FP8_PB_WO scale keys to SGLang block-FP8 params."""
+    if name.endswith(".weight_scale"):
+        candidate = name.removesuffix(".weight_scale") + ".weight_scale_inv"
+        if candidate in params_dict:
+            return candidate
+    return name
+
+
+def _get_k3_dense_weight(module: nn.Module) -> torch.Tensor:
+    """Return a dense weight with serialized block-FP8 scales applied."""
+    weight = module.weight.data
+    if not hasattr(module, "weight_scale_inv"):
+        return weight
+    return block_quant_dequant(
+        weight,
+        module.weight_scale_inv,
+        module.quant_method.weight_block_size,
+        module.params_dtype,
+    )
+
+
 def _k3_bf16_gemm(
     x: torch.Tensor,
     weight: torch.Tensor,
@@ -457,17 +487,17 @@ class KimiK3MoE(nn.Module):
             quant_config=quant_config,
             routed_scaling_factor=self.routed_scaling_factor,
             apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
-            # flashinfer_mxfp4 + situ consumes precomputed routing
+            # TRT-LLM cannot consume fused-front's row-strided router logits;
-            # (PackedPrecomputed): keep the radix router in the TopK layer
+            # keep K3's FP32 router and pass precomputed top-k instead.
-            # and hand its ids/weights to the MoE op. Other quantized paths
-            # keep the runner-resolved format (marlin -> standard anyway,
-            # bypassed only for the public logits-routing path).
             output_format=(
                 TopKOutputFormat.STANDARD
                 if quant_config is None
                 or (
                     config.hidden_act == "situ"
-                    and get_moe_runner_backend().is_flashinfer_mxfp4()
+                    and (
+                        get_moe_runner_backend().is_flashinfer_mxfp4()
+                        or get_moe_runner_backend().is_flashinfer_trtllm()
+                    )
                 )
                 # mega pre-dispatch consumes raw topk_ids/topk_weights
                 or get_moe_a2a_backend().is_megamoe()
@@ -1385,11 +1415,12 @@ class KimiK3DeltaAttention(nn.Module):
         self.use_full_rank_gate = config.linear_attn_config.get(
             "use_full_rank_gate", False
         )
+        self._bfa_uses_block_fp8 = self.use_full_rank_gate and _uses_modelopt_fp8_pb_wo(
+            quant_config, f"{prefix}.b_proj"
+        )
 
         # The fused path hardcodes tp_size sharding, so require attn_tp == tp.
-        # For the full-rank gate (K3) the checkpoint quantizes only the MoE
+        # Full-rank K3 also fuses mixed block-FP8 attention projections.
-        # experts; attention linears resolve to UnquantizedLinearMethod, so a
-        # non-None quant_config is fine for the merged projection.
         self.do_fuse_qkvbfg = self.attn_tp_size == self.tp_size and (
             quant_config is None or self.use_full_rank_gate
         )
@@ -1425,6 +1456,8 @@ class KimiK3DeltaAttention(nn.Module):
                 tp_rank=self.attn_tp_rank,
                 tp_size=self.attn_tp_size,
                 prefix=f"{prefix}.b_proj",
+                # TP8 shards K3's 96 beta rows below the 128-row FP8 block.
+                skip_block_quant_check=self._bfa_uses_block_fp8,
             )
             self.f_a_proj = ReplicatedLinear(
                 self.hidden_size,
@@ -1445,6 +1478,7 @@ class KimiK3DeltaAttention(nn.Module):
             # Merged [f_a | b] weight, built after weight loading by
             # _merge_bfa_weights().
             self._bfa_w: Optional[torch.Tensor] = None
+            self._bfa_f_b_w: Optional[torch.Tensor] = None
         elif self.do_fuse_qkvbfg:
             self.qkvb_sizes = [
                 projection_size,
@@ -1661,15 +1695,24 @@ class KimiK3DeltaAttention(nn.Module):
         and the width is padded to a multiple of 8 so every fused-output row
         stays 16-byte aligned for vectorized consumers (tiny-GEMM on f_b).
 
-        Called once from load_weights (after all weights are loaded, before
+        Called once after weight loading. Block-FP8 inputs are dequantized into
-        cuda graph capture)."""
+        the BF16 tiny-GEMM buffers here."""
         if not self.use_full_rank_gate:
             return
         if _is_npu:
             return
-        self._bfa_w, sizes = _merge_weights_as_views(
+        mods = [self.f_a_proj, self.b_proj]
-            [self.f_a_proj, self.b_proj], pad_rows_to=8
+        if self._bfa_uses_block_fp8:
-        )
+            weights = [_get_k3_dense_weight(mod) for mod in mods]
+            sizes = [weight.shape[0] for weight in weights]
+            pad = (-sum(sizes)) % 8
+            if pad:
+                weights.append(weights[0].new_zeros((pad, weights[0].shape[1])))
+            self._bfa_w = torch.cat(weights, dim=0).contiguous()
+            self._bfa_f_b_w = _get_k3_dense_weight(self.f_b_proj).contiguous()
+        else:
+            self._bfa_w, sizes = _merge_weights_as_views(mods, pad_rows_to=8)
+            self._bfa_f_b_w = self.f_b_proj.weight
         self._bfa_fa_size, self._bfa_b_size = sizes
 
     def _prepare_fused_decode(self) -> None:
@@ -1748,7 +1791,7 @@ class KimiK3DeltaAttention(nn.Module):
                     alt.wait_stream(cur)
                     with torch.cuda.stream(alt):
                         bfa = gemm(hidden_states, w)
-                        forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight)
+                        forget_gate = gemm(bfa[..., :n_fa], self._bfa_f_b_w)
                         beta = bfa[..., n_fa : n_fa + n_b]
                     fused_states, _ = self.fused_qkvg_proj(hidden_states)
                     qkv, g_proj_states = torch.split(
@@ -1760,7 +1803,7 @@ class KimiK3DeltaAttention(nn.Module):
                 fused_states, _ = self.fused_qkvg_proj(hidden_states)
                 qkv, g_proj_states = torch.split(fused_states, self.split_sizes, dim=-1)
                 bfa = gemm(hidden_states, w)
-                forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight)
+                forget_gate = gemm(bfa[..., :n_fa], self._bfa_f_b_w)
                 beta = bfa[..., n_fa : n_fa + n_b]
             else:
                 fused_states, _ = self.fused_qkvg_proj(hidden_states)
@@ -2883,6 +2926,8 @@ class KimiK3LinearForCausalLM(nn.Module):
         for args in weights:
             name, loaded_weight = args[:2]
             kwargs = args[2] if len(args) > 2 else {}
+            if name.endswith(".weight_scale") and loaded_weight.ndim == 4:
+                loaded_weight = loaded_weight[:, 0, :, 0]
 
             layer_id = get_layer_id(name)
             if layer_id is not None and (
@@ -2904,13 +2949,20 @@ class KimiK3LinearForCausalLM(nn.Module):
 
             # MLA: fuse q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj_with_mqa
             if ".q_a_proj." in name or ".kv_a_proj_with_mqa." in name:
+                is_q_a = ".q_a_proj." in name
                 fused_name = name.replace(".q_a_proj.", ".fused_qkv_a_proj_with_mqa.")
                 fused_name = fused_name.replace(
                     ".kv_a_proj_with_mqa.", ".fused_qkv_a_proj_with_mqa."
                 )
+                fused_name = _maybe_map_fp8_pb_scale_name(fused_name, params_dict)
                 if fused_name in params_dict:
                     param = params_dict[fused_name]
-                    if ".q_a_proj." in name:
+                    if fused_name.endswith(".weight_scale_inv"):
+                        offset = 0 if is_q_a else _cdiv(self.config.q_lora_rank, 128)
+                        param.data[offset : offset + loaded_weight.shape[0]].copy_(
+                            loaded_weight
+                        )
+                    elif is_q_a:
                         param.data[: loaded_weight.shape[0]].copy_(loaded_weight)
                     else:
                         q_lora_rank = self.config.q_lora_rank or 0
@@ -2947,6 +2999,7 @@ class KimiK3LinearForCausalLM(nn.Module):
                 name = name.replace(weight_name, param_name)
                 if name.endswith(".bias") and name not in params_dict:
                     continue
+                name = _maybe_map_fp8_pb_scale_name(name, params_dict)
                 param = params_dict[name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
@@ -2983,9 +3036,14 @@ class KimiK3LinearForCausalLM(nn.Module):
                     name = maybe_remap_kv_scale_name(name, params_dict)
                     if name is None:
                         continue
+                    name = _maybe_map_fp8_pb_scale_name(name, params_dict)
                     if name not in params_dict:
                         continue
                     param = params_dict[name]
+                    if name.endswith(".b_proj.weight_scale_inv"):
+                        # All TP ranks share K3's single beta output-scale block.
+                        param.data.copy_(loaded_weight)
+                    else:
                         weight_loader = getattr(
                             param, "weight_loader", default_weight_loader
                         )
@@ -3007,7 +3065,8 @@ class KimiK3LinearForCausalLM(nn.Module):
             if isinstance(layer, PPMissingLayer):
                 continue
             self_attn = layer.self_attn
-            w_kc, w_vc = self_attn.kv_b_proj.weight.unflatten(
+            kv_b_weight = _get_k3_dense_weight(self_attn.kv_b_proj)
+            w_kc, w_vc = kv_b_weight.unflatten(
                 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
             ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
             self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
@@ -3065,7 +3124,7 @@ class KimiK3LinearForCausalLM(nn.Module):
 
             if precompile_k3_recompute_w_u_kernel(
                 num_heads=layer.self_attn.local_num_heads,
-                dtype=layer.self_attn.o_proj.weight.dtype,
+                dtype=layer.self_attn.o_proj.params_dtype,
                 device=layer.self_attn.dt_bias.device,
             ):
                 rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.")

test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py

@@ -291,6 +291,26 @@ class TestG1ScaleC(CustomTestCase):
         self.assertTrue(g1_scale_c.is_contiguous())
         torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
 
+    def test_situ_keeps_both_dequant_scales_inside_activation(self):
+        # SiTU keeps both GEMM1 scales before tanh; scale_c contains only the
+        # GEMM2 input requantization.
+        num_experts = GATED_CONFIGS[0][1]
+        w2_input_scale_quant = torch.tensor(20.0)
+        gate = _global_scales(num_experts, 1, seed=8)
+        up = _global_scales(num_experts, 1, seed=9)
+
+        g1_scale_c = _compute_g1_scale_c(
+            w2_input_scale_quant,
+            gate,
+            up,
+            is_gated=True,
+            activation="situ",
+        )
+
+        self.assertEqual(g1_scale_c.shape, (num_experts,))
+        self.assertTrue(g1_scale_c.is_contiguous())
+        torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
+
     def test_scale_c_is_float32(self):
         # Lower-precision inputs are upcast to fp32 for the kernel.
         num_experts = GATED_CONFIGS[0][1]

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -696,6 +696,28 @@ class TestModelOptFp4LoaderSelection(CustomTestCase):
 
 
 class TestModelOptMixedPrecisionConfig(CustomTestCase):
+    def test_fp8_pb_wo_dispatches_to_native_block_fp8(self):
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"},
+                },
+                "packed_modules_mapping": {},
+            }
+        )
+
+        # Type dispatch only needs a LinearBase instance; skip GPU weight setup.
+        linear = ReplicatedLinear.__new__(ReplicatedLinear)
+        method = quant_config.get_quant_method(
+            linear, "model.layers.0.self_attn.q_proj"
+        )
+
+        self.assertIsInstance(method, Fp8LinearMethod)
+        self.assertEqual(method.quant_config.weight_block_size, [128, 128])
+        self.assertTrue(method.quant_config.is_checkpoint_fp8_serialized)
+        self.assertEqual(method.quant_config.activation_scheme, "dynamic")
+
     def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
         packed_modules_mapping = {
             "qkv_proj": ["q_proj", "k_proj", "v_proj"],

test/registered/unit/models/test_kimi_k3_bfa_overlap.py

@@ -8,7 +8,10 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.srt.models.kimi_k3 import KimiK3DeltaAttention
+from sglang.srt.models.kimi_k3 import (
+    KimiK3DeltaAttention,
+    _get_k3_dense_weight,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -39,9 +42,9 @@ def _make_owner(with_stream: bool):
     owner = SimpleNamespace(
         use_full_rank_gate=True,
         _bfa_w=_randn(_BFA_W_ROWS, _H).contiguous(),
+        _bfa_f_b_w=_randn(1536, _N_FA).contiguous(),
         _bfa_fa_size=_N_FA,
         _bfa_b_size=_N_B,
-        f_b_proj=SimpleNamespace(weight=_randn(1536, _N_FA).contiguous()),
         fused_qkvg_proj=fused_qkvg_proj,
         split_sizes=[3 * 1536, 1536],
         _bfa_alt_stream=torch.cuda.Stream() if with_stream else None,
@@ -97,6 +100,39 @@ class TestKimiK3BfaOverlap(CustomTestCase):
         for got, ref in zip(overlap, serial):
             self.assertTrue(torch.equal(got, ref))
 
+    def test_block_fp8_weight_is_dequantized_for_tiny_gemm(self):
+        module = SimpleNamespace(
+            weight=torch.nn.Parameter(
+                torch.ones((130, 129), device="cuda", dtype=torch.float8_e4m3fn),
+                requires_grad=False,
+            ),
+            weight_scale_inv=torch.nn.Parameter(
+                torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda"),
+                requires_grad=False,
+            ),
+            quant_method=SimpleNamespace(weight_block_size=[128, 128]),
+            params_dtype=torch.bfloat16,
+        )
+
+        weight = _get_k3_dense_weight(module)
+
+        self.assertEqual(weight.dtype, torch.bfloat16)
+        torch.testing.assert_close(
+            weight[[0, 0, 128, 128], [0, 128, 0, 128]].float(),
+            torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda"),
+        )
+
+    def test_per_tensor_fp8_weight_is_not_block_dequantized(self):
+        weight = torch.nn.Parameter(
+            torch.ones((2, 2), device="cuda", dtype=torch.float8_e4m3fn),
+            requires_grad=False,
+        )
+        module = SimpleNamespace(
+            weight=weight, weight_scale=torch.ones(1, device="cuda")
+        )
+
+        self.assertEqual(_get_k3_dense_weight(module).data_ptr(), weight.data_ptr())
+
 
 if __name__ == "__main__":
     unittest.main()