import os import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch from h3_blackwell_runtime.fc2_lt import ( _VALIDATED, _canonical_fc2_supported, _fc2_lt_compatible_result, _STATS, fc2_lt_enabled, fc2_lt_linear, prepare_fc2_lt, ) from h3_blackwell_runtime.nvfp4 import Nvfp4Linear from h3_blackwell_runtime.runtime import H3HotRuntime class Fc2LtContracts(unittest.TestCase): def setUp(self): _VALIDATED.clear() _fc2_lt_compatible_result.cache_clear() _STATS.update(attempts=0, successes=0, fallbacks=0) def test_flag_is_opt_in(self): with patch.dict(os.environ, {}, clear=True): self.assertFalse(fc2_lt_enabled()) with patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "yes"}, clear=True): self.assertTrue(fc2_lt_enabled()) def test_only_canonical_unsharded_fc2_is_eligible(self): linear = SimpleNamespace( role="h3_mlp_fc2", in_features=14_336, out_features=5_376, output_dtype=torch.bfloat16, bias=None, ) gate_up = SimpleNamespace( shape=(37_810, 28_672), device=torch.device("cuda", 0), ) qdata = SimpleNamespace(shape=(37_824, 7_168)) with patch("torch.cuda.get_device_capability", return_value=(12, 1)): self.assertTrue(_canonical_fc2_supported(linear, gate_up, qdata)) linear.in_features //= 2 self.assertFalse(_canonical_fc2_supported(linear, gate_up, qdata)) def test_prepare_loads_and_validates_extension(self): extension = MagicMock() extension.build_info.return_value = { "cuda_version": 13_000, "cublas_version": 130_100, "cuda_runtime_version": 13_000, "cublaslt_runtime_version": 130_000, } with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch( "h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None), ), ): self.assertTrue(prepare_fc2_lt()) extension.prepare.assert_called_once_with() def test_ineligible_shape_is_counted_as_fallback(self): args = (MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock()) with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=False), ): self.assertIsNone(fc2_lt_linear(*args)) self.assertEqual(_STATS["fallbacks"], 1) def test_h3_model_load_prepares_extension(self): model = MagicMock() model.eval.return_value = model runtime = SimpleNamespace( checkpoint=MagicMock(), config=SimpleNamespace(attention="sage2", mlp_chunks=1, mlp_chunk_threshold=4096), ) with ( patch( "h3_blackwell_runtime.runtime.H3PackedDenoiser.from_checkpoint", return_value=model, ), patch("h3_blackwell_runtime.runtime.configure_mlp_chunking") as configure, patch("h3_blackwell_runtime.fc2_lt.prepare_fc2_lt") as prepare, ): self.assertIs(H3HotRuntime._load_h3(runtime), model) configure.assert_called_once_with(model, 1, 4096) prepare.assert_called_once_with() def test_extension_load_failure_falls_back_unless_strict(self): args = (MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock()) with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(None, RuntimeError("load"))), ): self.assertIsNone(fc2_lt_linear(*args)) with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1", "H3_NVFP4_FC2_LT_STRICT": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(None, RuntimeError("load"))), ): with self.assertRaisesRegex(RuntimeError, "failed to load"): fc2_lt_linear(*args) def test_build_info_and_runtime_version_failures_fall_back(self): extension = MagicMock() extension.build_info.side_effect = RuntimeError("metadata") args = (MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock()) with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None)), ): self.assertIsNone(fc2_lt_linear(*args)) _fc2_lt_compatible_result.cache_clear() extension.build_info.side_effect = None extension.build_info.return_value = { "cuda_version": 13_000, "cublas_version": 130_100, "cuda_runtime_version": 12_900, "cublaslt_runtime_version": 130_000, } with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None)), ): self.assertIsNone(fc2_lt_linear(*args)) def test_algo_check_and_execution_failures_fall_back(self): extension = MagicMock() extension.build_info.return_value = { "cuda_version": 13_000, "cublas_version": 130_100, "cuda_runtime_version": 13_000, "cublaslt_runtime_version": 130_000, } linear = MagicMock() linear.weight.shape = (5_376, 7_168) gate_up = MagicMock() gate_up.device.index = 0 qdata = MagicMock() qdata.shape = (37_824, 7_168) args = (linear, gate_up, MagicMock(), qdata, MagicMock()) common = ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None)), ) extension.check.return_value = {"valid": False, "required_workspace_bytes": 0} with common[0], common[1], common[2]: self.assertIsNone(fc2_lt_linear(*args)) _VALIDATED.clear() extension.check.return_value = {"valid": True, "required_workspace_bytes": 0} extension.run.side_effect = RuntimeError("run") with ( patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch("h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None)), patch("h3_blackwell_runtime.fc2_lt.torch.zeros_like", return_value=MagicMock()), patch("h3_blackwell_runtime.fc2_lt.torch.empty", return_value=MagicMock()), ): self.assertIsNone(fc2_lt_linear(*args)) def test_strict_mode_raises_on_algo_check_and_execution_failures(self): extension = MagicMock() extension.build_info.return_value = { "cuda_version": 13_000, "cublas_version": 130_100, "cuda_runtime_version": 13_000, "cublaslt_runtime_version": 130_000, } linear = MagicMock() linear.weight.shape = (5_376, 7_168) gate_up = MagicMock() gate_up.device.index = 0 qdata = MagicMock() qdata.shape = (37_824, 7_168) args = (linear, gate_up, MagicMock(), qdata, MagicMock()) environment = { "H3_NVFP4_FC2_LT_SPLITK1": "1", "H3_NVFP4_FC2_LT_STRICT": "1", } extension.check.return_value = {"valid": False, "required_workspace_bytes": 0} with ( patch.dict(os.environ, environment, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch( "h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None), ), ): with self.assertRaisesRegex(RuntimeError, "algorithm check rejected"): fc2_lt_linear(*args) _VALIDATED.clear() _fc2_lt_compatible_result.cache_clear() extension.check.return_value = {"valid": True, "required_workspace_bytes": 0} extension.run.side_effect = RuntimeError("run") with ( patch.dict(os.environ, environment, clear=True), patch("h3_blackwell_runtime.fc2_lt._canonical_fc2_supported", return_value=True), patch( "h3_blackwell_runtime.fc2_lt._fc2_lt_extension_result", return_value=(extension, None), ), patch("h3_blackwell_runtime.fc2_lt.torch.zeros_like", return_value=MagicMock()), patch("h3_blackwell_runtime.fc2_lt.torch.empty", return_value=MagicMock()), ): with self.assertRaisesRegex(RuntimeError, "execution failed"): fc2_lt_linear(*args) def test_forward_swiglu_uses_schedule_without_repacking(self): linear = SimpleNamespace( role="h3_mlp_fc2", in_features=4, out_features=3, output_dtype=torch.bfloat16, bias=None, pre_quant_scale=None, active_lora=None, lora_strength=0.0, full_precision_matrix_mult=False, ) gate_up = torch.randn(1, 8, dtype=torch.bfloat16) packed = (MagicMock(), MagicMock(), MagicMock()) expected = torch.randn(1, 3, dtype=torch.bfloat16) with ( torch.inference_mode(), patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch.object(torch.Tensor, "is_cuda", new_callable=unittest.mock.PropertyMock, return_value=True), patch("h3_blackwell_runtime.nvfp4_quant.vortex_native_quantize_swiglu_nvfp4", return_value=packed) as quantize, patch("h3_blackwell_runtime.nvfp4_quant.wrap_native_swiglu_nvfp4") as wrap, patch("h3_blackwell_runtime.fc2_lt.fc2_lt_linear", return_value=expected) as scheduled, ): actual = Nvfp4Linear.forward_swiglu(linear, gate_up) self.assertTrue(torch.equal(actual, expected)) quantize.assert_called_once_with(gate_up) scheduled.assert_called_once_with(linear, gate_up, *packed) wrap.assert_not_called() def test_forward_swiglu_reuses_native_bytes_on_fallback(self): packed_weight = MagicMock() linear = SimpleNamespace( role="h3_mlp_fc2", in_features=4, out_features=3, output_dtype=torch.bfloat16, bias=None, pre_quant_scale=None, active_lora=None, lora_strength=0.0, full_precision_matrix_mult=False, _packed_weight=lambda: packed_weight, ) gate_up = torch.randn(1, 8, dtype=torch.bfloat16) packed = (MagicMock(), MagicMock(), MagicMock()) wrapped = MagicMock() expected = torch.randn(1, 3, dtype=torch.bfloat16) with ( torch.inference_mode(), patch.dict(os.environ, {"H3_NVFP4_FC2_LT_SPLITK1": "1"}, clear=True), patch.object(torch.Tensor, "is_cuda", new_callable=unittest.mock.PropertyMock, return_value=True), patch("h3_blackwell_runtime.nvfp4_quant.vortex_native_quantize_swiglu_nvfp4", return_value=packed) as quantize, patch("h3_blackwell_runtime.nvfp4_quant.wrap_native_swiglu_nvfp4", return_value=wrapped) as wrap, patch("h3_blackwell_runtime.fc2_lt.fc2_lt_linear", return_value=None), patch("h3_blackwell_runtime.nvfp4.functional.linear", return_value=expected) as gemm, ): actual = Nvfp4Linear.forward_swiglu(linear, gate_up) self.assertTrue(torch.equal(actual, expected)) quantize.assert_called_once_with(gate_up) wrap.assert_called_once_with(gate_up, packed) gemm.assert_called_once_with(wrapped, packed_weight, None) if __name__ == "__main__": unittest.main()