287 lines
12 KiB
Python
287 lines
12 KiB
Python
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()
|