From d903351a669dcda7742821af83ebc0ce99f674f5 Mon Sep 17 00:00:00 2001 From: AndyLi429 <68410213+AndyLi429@users.noreply.github.com> Date: Sun, 20 Sep 2026 09:55:30 +0800 Subject: [PATCH] [NPU][bugfix] update low latency quantization input and update MXFP8 tests (#38831) Co-authored-by: AndyLi429 --- .../srt/layers/moe/token_dispatcher/deepep.py | 155 ++++--- python/sglang/srt/layers/moe/utils.py | 2 + .../npu/quantization/test_fp4_moe_methods.py | 430 ++++++++++++++++++ 3 files changed, 518 insertions(+), 69 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 4502751cb..6e7588d9e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -441,10 +441,14 @@ class _DeepEPDispatcherImplBase: config_map = { DispatcherOutputDtype.BF16: { "use_fp8": False, + "use_mxfp4": False, + "use_mxfp8": False, "use_nvfp4": False, }, DispatcherOutputDtype.FP8: { "use_fp8": True, + "use_mxfp4": False, + "use_mxfp8": False, "use_nvfp4": False, }, # Needed for Ascend A2/A3 NPU case, @@ -452,14 +456,26 @@ class _DeepEPDispatcherImplBase: # quantization will be performed in int8 DispatcherOutputDtype.INT8: { "use_fp8": True, + "use_mxfp4": False, + "use_mxfp8": False, "use_nvfp4": False, }, DispatcherOutputDtype.NVFP4: { "use_fp8": False, + "use_mxfp4": False, + "use_mxfp8": False, "use_nvfp4": True, }, + DispatcherOutputDtype.MXFP4: { + "use_fp8": False, + "use_mxfp4": True, + "use_mxfp8": False, + "use_nvfp4": False, + }, DispatcherOutputDtype.MXFP8: { "use_fp8": False, + "use_mxfp4": False, + "use_mxfp8": True, "use_nvfp4": False, }, } @@ -470,6 +486,8 @@ class _DeepEPDispatcherImplBase: # Apply configuration config = config_map[self.deepep_output_dtype] self.use_fp8 = config["use_fp8"] + self.use_mxfp4 = config["use_mxfp4"] + self.use_mxfp8 = config["use_mxfp8"] self.use_nvfp4 = config["use_nvfp4"] # Handle environment variables @@ -478,33 +496,17 @@ class _DeepEPDispatcherImplBase: def _validate_and_adjust_dtype(self) -> None: """Validate dtype against hardware and adjust if necessary.""" - self.low_latency_quant_mode = None - self._low_latency_quant_mode_runtime_checked = False - if self.deepep_output_dtype == DispatcherOutputDtype.MXFP8: - if not _is_npu or self.dispatch_mode != DeepEPMode.LOW_LATENCY: - raise RuntimeError( - "MXFP8 DeepEP dispatch is supported only for A5 " - "low-latency dispatch." - ) - + if _is_npu and self.deepep_output_dtype == DispatcherOutputDtype.FP8: from sglang.srt.hardware_backend.npu.utils import is_npu_arch35 if not is_npu_arch35(): - raise RuntimeError( - "MXFP8 DeepEP dispatch is supported only on Ascend A5 " - "in low-latency mode." - ) - self.low_latency_quant_mode = "mx_fp8_e4m3" - return - - if _is_npu: - if self.deepep_output_dtype == DispatcherOutputDtype.FP8: logger.warning_once( "Ascend A2/A3 NPU does not support fp8 " - "deepep_dispatcher_output_dtype, switching to int8..." + "deepep_dispatcher_output_dtype; DeepEP will use int8." ) - self.deepep_output_dtype = DispatcherOutputDtype.INT8 - elif self.deepep_output_dtype == DispatcherOutputDtype.NVFP4: + + if _is_npu: + if self.deepep_output_dtype == DispatcherOutputDtype.NVFP4: raise RuntimeError( "Ascend A2/A3 NPU does not support nvfp4 deepep_dispatcher_output_dtype." ) @@ -561,6 +563,40 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): previous_event = Buffer.capture() if self.async_finish else None return hidden_states, topk_ids, topk_weights, previous_event + def _get_quantization_kwargs(self, buffer: Buffer) -> dict: + if not _is_npu: + return {} + + dispatch_params = inspect.signature(buffer.dispatch).parameters + flag_kwargs = { + "use_fp8": self.use_fp8, + "use_mxfp4": self.use_mxfp4, + "use_mxfp8": self.use_mxfp8, + } + if all(name in dispatch_params for name in flag_kwargs): + return flag_kwargs + + if "quant_mode" in dispatch_params: + if self.use_mxfp4: + quant_mode = "mx_fp4_e2m1" + elif self.use_mxfp8: + quant_mode = "mx_fp8_e4m3" + elif self.use_fp8: + quant_mode = "int8" + else: + quant_mode = "bf16" + return {"quant_mode": quant_mode} + + if not self.use_mxfp4 and not self.use_mxfp8: + # A3's legacy pybind Buffer does not expose its dispatch signature. + # It selects BF16/INT8 dispatch through the DeepEP runtime instead. + return {} + + raise RuntimeError( + "Installed DeepEP normal dispatch does not support either " + "use_fp8/use_mxfp4/use_mxfp8 or quant_mode." + ) + def dispatch_b(self, hidden_states, topk_ids, topk_weights, previous_event): ( hidden_states, @@ -610,6 +646,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): # `handle` as a member variable works. _deepep_precompile_tp_barrier() + npu_quantization_opts = self._get_quantization_kwargs(buffer) ( recv_x, recv_topk_ids, @@ -630,6 +667,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): allocate_on_comm_stream=(previous_event is not None) and self.async_finish, expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1, config=DeepEPConfig.get_instance().normal_dispatch_config, + **npu_quantization_opts, ) get_global_expert_distribution_recorder().on_deepep_dispatch_normal( num_recv_tokens_per_expert, @@ -774,9 +812,8 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): ): input_global_scale = self.quant_config.get("input_global_scale", None) - # round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP - # to return int32-packed UE8M0 scales that don't feed the flashinfer - # cutedsl kernel. + # round_scale / use_ue8m0 are FP8-DeepGEMM specific. Dropping use_ue8m0 + # makes DeepEP return fp32 column-major scales the e8m0 cast cannot view. fp8_deepgemm_scale_opts = ( dict( round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM @@ -789,58 +826,16 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): ) buffer = self._get_buffer() - if ( - self.low_latency_quant_mode is not None - and not self._low_latency_quant_mode_runtime_checked - ): - try: - dispatch_signature = inspect.signature(buffer.low_latency_dispatch) - except (TypeError, ValueError) as exc: - raise RuntimeError( - "A5 MXFP8 DeepEP dispatch requires a recent " - "sgl-kernel-npu/DeepEP runtime exposing " - "low_latency_dispatch(..., quant_mode=...)." - ) from exc - if "quant_mode" not in dispatch_signature.parameters: - raise RuntimeError( - "A5 MXFP8 DeepEP dispatch requires a recent " - "sgl-kernel-npu/DeepEP runtime exposing " - "low_latency_dispatch(..., quant_mode=...)." - ) - self._low_latency_quant_mode_runtime_checked = True - - use_fp8 = self.use_fp8 - low_latency_quant_kwargs = {} - if self.low_latency_quant_mode is not None: - deep_use_mode = os.environ.get("DEEP_USE_MODE", "default") - if deep_use_mode == "default": - low_latency_quant_kwargs = { - "quant_mode": self.low_latency_quant_mode, - } - elif deep_use_mode == "ops": - # The ops strategy ignores quant_mode and uses the legacy - # flags. Pass both forms so the request is explicit and the - # strategy still produces E4M3 + E8M0 MXFP8 tensors. - use_fp8 = True - low_latency_quant_kwargs = { - "quant_mode": self.low_latency_quant_mode, - "use_ue8m0": True, - } - else: - raise RuntimeError( - "A5 MXFP8 DeepEP dispatch supports only " - "DEEP_USE_MODE=default or DEEP_USE_MODE=ops; got " - f"{deep_use_mode!r}." - ) _deepep_precompile_tp_barrier() + npu_mxfp_quantization_opts = self._get_npu_mxfp_quantization_kwargs(buffer) packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( buffer.low_latency_dispatch( hidden_states, topk_ids, self.num_max_dispatch_tokens_per_rank, self.num_experts, - use_fp8=use_fp8, - **low_latency_quant_kwargs, + use_fp8=self.use_fp8, + **npu_mxfp_quantization_opts, **( dict(topk_weights=topk_weights) if _is_npu and not _use_zbal @@ -859,6 +854,28 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): ) return packed_recv_hidden, self.packed_recv_count, event, hook + def _get_npu_mxfp_quantization_kwargs(self, buffer: Buffer) -> dict: + if not _is_npu: + return {} + + parameters = inspect.signature(buffer.low_latency_dispatch).parameters + if any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ): + return { + "use_mxfp4": self.use_mxfp4, + "use_mxfp8": self.use_mxfp8, + } + return { + name: value + for name, value in { + "use_mxfp4": self.use_mxfp4, + "use_mxfp8": self.use_mxfp8, + }.items() + if name in parameters + } + def combine_a( self, hidden_states: torch.Tensor, diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index d6c3ea032..9840e671f 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -284,6 +284,7 @@ class DispatcherOutputDtype(Enum): - FP8: dispatch hidden states in fp8 - INT8: dispatch hidden states in int8 - NVFP4: dispatch hidden states in nvfp4 + - MXFP4: dispatch hidden states in mxfp4 (fp4_e2m1 + e8m0 block scale) - MXFP8: dispatch hidden states in mxfp8 (fp8_e4m3 + e8m0 block scale) """ @@ -291,6 +292,7 @@ class DispatcherOutputDtype(Enum): FP8 = "fp8" INT8 = "int8" NVFP4 = "nvfp4" + MXFP4 = "mxfp4" MXFP8 = "mxfp8" diff --git a/test/registered/unit/npu/quantization/test_fp4_moe_methods.py b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py index 1594ae8b8..a2797a6df 100644 --- a/test/registered/unit/npu/quantization/test_fp4_moe_methods.py +++ b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py @@ -27,6 +27,7 @@ from sglang.srt.hardware_backend.npu.quantization.moe_methods import ( w4a8_mxfp_gmm, ) from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.moe.token_dispatcher import deepep from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod @@ -186,6 +187,435 @@ class TestPairPackMxfpActScale(unittest.TestCase): _pair_pack_mxfp_act_scale(torch.zeros(2, 3)) +class _LowLatencyBuffer: + """The MXFP8-era Buffer: bool flags, no quant_mode.""" + + def __init__(self): + self.kwargs = None + + def low_latency_dispatch( + self, + hidden_states, + topk_ids, + num_max_dispatch_tokens_per_rank, + num_experts, + *, + use_fp8, + use_mxfp4=False, + use_mxfp8=False, + **kwargs, + ): + self.kwargs = { + "use_fp8": use_fp8, + "use_mxfp4": use_mxfp4, + "use_mxfp8": use_mxfp8, + **kwargs, + } + return torch.empty(0), torch.empty(0), object(), object(), object() + + +class _LegacyLowLatencyBuffer: + """A DeepEP API version that predates the use_mxfp8 flag.""" + + def __init__(self): + self.use_mxfp4 = None + + def low_latency_dispatch( + self, + hidden_states, + topk_ids, + num_max_dispatch_tokens_per_rank, + num_experts, + *, + use_fp8, + use_mxfp4=False, + topk_weights, + async_finish, + return_recv_hook, + ): + self.use_mxfp4 = use_mxfp4 + return torch.empty(0), torch.empty(0), object(), object(), object() + + +class _CudaLowLatencyBuffer: + """CUDA's Buffer API does not accept the NPU-only MXFP flags.""" + + def __init__(self): + self.use_fp8 = None + + def low_latency_dispatch( + self, + hidden_states, + topk_ids, + num_max_dispatch_tokens_per_rank, + num_experts, + *, + use_fp8, + round_scale=False, + use_ue8m0=False, + async_finish=False, + return_recv_hook=False, + ): + self.use_fp8 = use_fp8 + return torch.empty(0), torch.empty(0), object(), object(), object() + + +class _CudaNormalBuffer: + """CUDA's normal Buffer API receives an already-quantized input tuple.""" + + def __init__(self): + self.dispatched = False + + def get_dispatch_layout(self, *args, **kwargs): + return ( + torch.ones(1, dtype=torch.int32), + None, + torch.ones(2, dtype=torch.int32), + torch.ones(1, 1, dtype=torch.bool), + None, + ) + + def dispatch( + self, + x, + *, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert, + previous_event, + async_finish, + allocate_on_comm_stream, + expert_alignment, + config, + ): + self.dispatched = True + return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object() + + +class _FlagNormalBuffer(_CudaNormalBuffer): + def __init__(self): + super().__init__() + self.quantization_kwargs = None + + def dispatch( + self, + x, + *, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert, + previous_event, + async_finish, + allocate_on_comm_stream, + expert_alignment, + config, + use_fp8, + use_mxfp4, + use_mxfp8, + ): + self.dispatched = True + self.quantization_kwargs = { + "use_fp8": use_fp8, + "use_mxfp4": use_mxfp4, + "use_mxfp8": use_mxfp8, + } + return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object() + + +class _LegacyNormalBuffer: + """The pre-bool-flags DeepEP normal-dispatch API used by CI.""" + + def __init__(self): + self.quant_mode = None + + def get_dispatch_layout(self, *args, **kwargs): + return ( + torch.ones(1, dtype=torch.int32), + None, + torch.ones(2, dtype=torch.int32), + torch.ones(1, 1, dtype=torch.bool), + None, + ) + + def dispatch( + self, + x, + *, + topk_idx, + topk_weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert, + previous_event, + async_finish, + allocate_on_comm_stream, + expert_alignment, + config, + quant_mode, + ): + self.quant_mode = quant_mode + return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object() + + +class _OpaqueNormalBuffer(_CudaNormalBuffer): + """The A3 pybind Buffer API whose dispatch signature hides quantization args.""" + + def __init__(self): + super().__init__() + self.dispatch_kwargs = None + + def dispatch(self, *args, **kwargs): + self.dispatched = True + self.dispatch_kwargs = kwargs + return torch.empty(0), torch.empty(0), torch.empty(0), [], object(), object() + + +class TestDeepEPLowLatencyMxfp8Dispatch(unittest.TestCase): + def test_mxfp4_output_dtype_enables_only_mxfp4(self): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplBase) + + with patch.object( + deepep, + "get_deepep_output_dtype", + return_value=deepep.DispatcherOutputDtype.MXFP4, + ): + dispatcher.set_deepep_dispatcher_dtype() + + self.assertFalse(dispatcher.use_fp8) + self.assertTrue(dispatcher.use_mxfp4) + self.assertFalse(dispatcher.use_mxfp8) + + @staticmethod + def _dispatcher(quant_mode, buffer): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplLowLatency) + dispatcher.quant_config = {} + dispatcher.use_fp8 = False + dispatcher.use_mxfp4 = False + dispatcher.use_mxfp8 = quant_mode == "mxfp8" + dispatcher.use_nvfp4 = False + dispatcher.num_max_dispatch_tokens_per_rank = 2 + dispatcher.num_experts = 2 + dispatcher.return_recv_hook = False + dispatcher._get_buffer = lambda: buffer + return dispatcher + + def test_mxfp8_passes_the_mxfp8_flag_without_ue8m0(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher("mxfp8", buffer) + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertFalse(buffer.kwargs["use_fp8"]) + self.assertTrue(buffer.kwargs["use_mxfp8"]) + self.assertNotIn("use_ue8m0", buffer.kwargs) + + def test_mxfp4_passes_the_mxfp4_flag(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher("mxfp4", buffer) + dispatcher.use_mxfp4 = True + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertFalse(buffer.kwargs["use_fp8"]) + self.assertTrue(buffer.kwargs["use_mxfp4"]) + self.assertFalse(buffer.kwargs["use_mxfp8"]) + + def test_bf16_omits_unsupported_mxfp8_flag_for_legacy_buffer(self): + buffer = _LegacyLowLatencyBuffer() + dispatcher = self._dispatcher("bf16", buffer) + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertFalse(buffer.use_mxfp4) + + def test_normal_dispatch_passes_quantization_flags(self): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal) + dispatcher.num_experts = 2 + dispatcher.async_finish = False + dispatcher.use_fp8 = False + dispatcher.use_mxfp4 = False + dispatcher.use_mxfp8 = True + buffer = _FlagNormalBuffer() + dispatcher._get_buffer = lambda: buffer + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + patch.object( + deepep.DeepEPConfig, + "get_instance", + return_value=SimpleNamespace(normal_dispatch_config=None), + ), + patch.object( + deepep, + "get_global_expert_distribution_recorder", + return_value=MagicMock(), + ), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + None, + ) + + self.assertTrue(buffer.quantization_kwargs["use_mxfp8"]) + self.assertFalse(buffer.quantization_kwargs["use_fp8"]) + self.assertFalse(buffer.quantization_kwargs["use_mxfp4"]) + + def test_normal_dispatch_uses_legacy_quant_mode_when_flags_are_unsupported(self): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal) + dispatcher.num_experts = 2 + dispatcher.async_finish = False + dispatcher.use_fp8 = False + dispatcher.use_mxfp4 = False + dispatcher.use_mxfp8 = False + buffer = _LegacyNormalBuffer() + dispatcher._get_buffer = lambda: buffer + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + patch.object( + deepep.DeepEPConfig, + "get_instance", + return_value=SimpleNamespace(normal_dispatch_config=None), + ), + patch.object( + deepep, + "get_global_expert_distribution_recorder", + return_value=MagicMock(), + ), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + None, + ) + + self.assertEqual(buffer.quant_mode, "bf16") + + def test_normal_dispatch_keeps_a3_legacy_path_for_opaque_signature(self): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal) + dispatcher.num_experts = 2 + dispatcher.async_finish = False + dispatcher.use_fp8 = True + dispatcher.use_mxfp4 = False + dispatcher.use_mxfp8 = False + buffer = _OpaqueNormalBuffer() + dispatcher._get_buffer = lambda: buffer + + with ( + patch.object(deepep, "_is_npu", True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + patch.object( + deepep.DeepEPConfig, + "get_instance", + return_value=SimpleNamespace(normal_dispatch_config=None), + ), + patch.object( + deepep, + "get_global_expert_distribution_recorder", + return_value=MagicMock(), + ), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + None, + ) + + self.assertTrue(buffer.dispatched) + self.assertNotIn("use_fp8", buffer.dispatch_kwargs) + self.assertNotIn("use_mxfp4", buffer.dispatch_kwargs) + self.assertNotIn("use_mxfp8", buffer.dispatch_kwargs) + self.assertNotIn("quant_mode", buffer.dispatch_kwargs) + + def test_cuda_normal_dispatch_omits_npu_quantization_flags(self): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplNormal) + dispatcher.num_experts = 2 + dispatcher.async_finish = False + dispatcher.use_fp8 = True + dispatcher.use_mxfp4 = True + dispatcher.use_mxfp8 = True + buffer = _CudaNormalBuffer() + dispatcher._get_buffer = lambda: buffer + + with ( + patch.object(deepep, "_is_npu", False), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + patch.object( + deepep.DeepEPConfig, + "get_instance", + return_value=SimpleNamespace(normal_dispatch_config=None), + ), + patch.object( + deepep, + "get_global_expert_distribution_recorder", + return_value=MagicMock(), + ), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + None, + ) + + self.assertTrue(buffer.dispatched) + + def test_cuda_low_latency_dispatch_omits_npu_mxfp_flags(self): + buffer = _CudaLowLatencyBuffer() + dispatcher = self._dispatcher("mxfp8", buffer) + dispatcher.use_fp8 = True + dispatcher.use_mxfp4 = True + + with ( + patch.object(deepep, "_is_npu", False), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertTrue(buffer.use_fp8) + + class TestW4A8MxfpGmmInputScale(unittest.TestCase): def setUp(self): self.input = torch.randn(2, 64)