From 5f128395910dafb98c34083dc26cb790c7674d34 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:13:45 -0700 Subject: [PATCH] [Fix] Support Kimi-K3 ModelOpt mixed NVFP4/FP8 checkpoint (#35077) --- .../moe/moe_runner/flashinfer_trtllm.py | 15 ++- .../srt/layers/quantization/modelopt_quant.py | 41 +++++-- python/sglang/srt/models/kimi_k3.py | 105 ++++++++++++++---- .../test_modelopt_nvfp4_moe_scales.py | 20 ++++ .../unit/model_loader/test_modelopt_loader.py | 22 ++++ .../unit/models/test_kimi_k3_bfa_overlap.py | 40 ++++++- 6 files changed, 205 insertions(+), 38 deletions(-) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index a8bf5668d..9e72c3c04 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/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}'." diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index a521b04a2..96fbd2819 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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 - or layer.moe_runner_config.swiglu_limit + None + 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( - layer, - "gemm1_beta", - (1.0 / layer.g1_alphas).to(torch.float32), + gemm1_beta = ( + torch.full_like( + layer.g1_alphas, + 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 ( - activation in _SUPPORTED_ACT_STRS - ), f"{activation=} not in supported {_SUPPORTED_ACT_STRS}" + assert activation in _SUPPORTED_ACT_STRS or ( + activation == "situ" and moe_runner_backend.is_flashinfer_trtllm() + ), f"{activation=} is unsupported by {moe_runner_backend}" moe_runner_config = self.moe_runner_config if moe_runner_backend.is_marlin(): diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 1dfb8f0cc..f910195d0 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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 - # (PackedPrecomputed): keep the radix router in the TopK layer - # 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). + # TRT-LLM cannot consume fused-front's row-strided router logits; + # keep K3's FP32 router and pass precomputed top-k instead. 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 - # experts; attention linears resolve to UnquantizedLinearMethod, so a - # non-None quant_config is fine for the merged projection. + # Full-rank K3 also fuses mixed block-FP8 attention projections. 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 - cuda graph capture).""" + Called once after weight loading. Block-FP8 inputs are dequantized into + 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( - [self.f_a_proj, self.b_proj], pad_rows_to=8 - ) + mods = [self.f_a_proj, self.b_proj] + 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,13 +3036,18 @@ 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] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight, **kwargs) + 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 + ) + weight_loader(param, loaded_weight, **kwargs) loaded_params.add(name) self.post_load_weights() @@ -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.") diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py index fedcc786c..5424212f7 100644 --- a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py +++ b/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] diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index eddc599f2..ea642a0eb 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/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"], diff --git a/test/registered/unit/models/test_kimi_k3_bfa_overlap.py b/test/registered/unit/models/test_kimi_k3_bfa_overlap.py index 128fc28c2..db856e293 100644 --- a/test/registered/unit/models/test_kimi_k3_bfa_overlap.py +++ b/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()