123 lines
5.7 KiB
Python
123 lines
5.7 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
|
|
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
|
|
class Nvfp4NativePackParity(unittest.TestCase):
|
|
def test_h3_projection_widths_match_comfy_layout(self):
|
|
try:
|
|
import comfy_kitchen as ck
|
|
from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
|
|
from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
|
|
except ImportError as error:
|
|
raise unittest.SkipTest("Comfy Kitchen is not installed") from error
|
|
|
|
torch.manual_seed(440420)
|
|
extension = _vortex_scale_extension()
|
|
for features in (5376, 7168, 14336):
|
|
with self.subTest(features=features):
|
|
value = torch.randn(129, features, device="cuda", dtype=torch.bfloat16)
|
|
scale = (value.abs().amax() / float(F8_E4M3_MAX * F4_E2M1_MAX)).float()
|
|
expected_qdata, expected_block_scale = ck.quantize_nvfp4(value, scale, pad_16x=True)
|
|
actual_qdata, actual_block_scale = extension.quantize_nvfp4_bf16(value, scale, True, 512)
|
|
self.assertTrue(torch.equal(actual_qdata, expected_qdata))
|
|
self.assertTrue(torch.equal(actual_block_scale.view(torch.uint8), expected_block_scale.view(torch.uint8)))
|
|
into_qdata = torch.empty_like(actual_qdata)
|
|
into_block_scale = torch.empty_like(actual_block_scale)
|
|
extension.quantize_nvfp4_bf16_into(
|
|
value, scale, into_qdata, into_block_scale, 512,
|
|
)
|
|
self.assertTrue(torch.equal(into_qdata, expected_qdata))
|
|
self.assertTrue(
|
|
torch.equal(
|
|
into_block_scale.view(torch.uint8),
|
|
expected_block_scale.view(torch.uint8),
|
|
)
|
|
)
|
|
low_first_qdata = torch.empty_like(actual_qdata)
|
|
extension.quantize_nvfp4_bf16_into(
|
|
value,
|
|
scale,
|
|
low_first_qdata,
|
|
into_block_scale,
|
|
512,
|
|
False,
|
|
)
|
|
expected_low_first = (
|
|
((expected_qdata & 0x0F) << 4)
|
|
| ((expected_qdata & 0xF0) >> 4)
|
|
)
|
|
self.assertTrue(torch.equal(low_first_qdata, expected_low_first))
|
|
|
|
def test_modulated_producer_matches_materialized_bf16(self):
|
|
try:
|
|
import comfy_kitchen as ck
|
|
from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
|
|
from h3_blackwell_runtime.h3_fusion import fused_modulate_
|
|
from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
|
|
except ImportError as error:
|
|
raise unittest.SkipTest("Comfy Kitchen or Triton is not installed") from error
|
|
|
|
torch.manual_seed(440421)
|
|
value = torch.randn(129, 128, device="cuda", dtype=torch.bfloat16)
|
|
row_index = torch.arange(129, device="cuda", dtype=torch.int32) % 3
|
|
divisor = float(F8_E4M3_MAX * F4_E2M1_MAX)
|
|
extension = _vortex_scale_extension()
|
|
for table_dtype in (torch.bfloat16, torch.float32):
|
|
with self.subTest(table_dtype=table_dtype):
|
|
shift = torch.randn(3, 128, device="cuda", dtype=table_dtype)
|
|
scale = torch.randn(3, 128, device="cuda", dtype=table_dtype)
|
|
materialized = fused_modulate_(
|
|
value.clone(), shift, scale, row_index,
|
|
)
|
|
expected_scale = (materialized.abs().amax() / divisor).to(
|
|
torch.bfloat16,
|
|
).float()
|
|
expected_qdata, expected_sfa = ck.quantize_nvfp4(
|
|
materialized, expected_scale, pad_16x=True,
|
|
)
|
|
actual_scale, actual_qdata, actual_sfa = (
|
|
extension.quantize_nvfp4_modulated_bf16(
|
|
value, shift, scale, row_index, divisor, True, 32, 128, 128,
|
|
)
|
|
)
|
|
self.assertTrue(torch.equal(actual_scale, expected_scale))
|
|
self.assertTrue(torch.equal(actual_qdata, expected_qdata))
|
|
self.assertTrue(
|
|
torch.equal(actual_sfa.view(torch.uint8), expected_sfa.view(torch.uint8))
|
|
)
|
|
|
|
def test_swiglu_producer_matches_materialized_bf16(self):
|
|
try:
|
|
import comfy_kitchen as ck
|
|
from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
|
|
from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
|
|
except ImportError as error:
|
|
raise unittest.SkipTest("Comfy Kitchen is not installed") from error
|
|
|
|
torch.manual_seed(440422)
|
|
gate_up = torch.randn(129, 256, device="cuda", dtype=torch.bfloat16)
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
materialized = torch.nn.functional.silu(gate).mul_(up)
|
|
divisor = float(F8_E4M3_MAX * F4_E2M1_MAX)
|
|
expected_scale = (materialized.abs().amax() / divisor).to(
|
|
torch.bfloat16,
|
|
).float()
|
|
expected_qdata, expected_sfa = ck.quantize_nvfp4(
|
|
materialized, expected_scale, pad_16x=True,
|
|
)
|
|
actual_scale, actual_qdata, actual_sfa = (
|
|
_vortex_scale_extension().quantize_nvfp4_swiglu_bf16(
|
|
gate_up, divisor, True, 32, 128, 128,
|
|
)
|
|
)
|
|
self.assertTrue(torch.equal(actual_scale, expected_scale))
|
|
self.assertTrue(torch.equal(actual_qdata, expected_qdata))
|
|
self.assertTrue(
|
|
torch.equal(actual_sfa.view(torch.uint8), expected_sfa.view(torch.uint8))
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|