67 lines
2.8 KiB
Python
67 lines
2.8 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.block import gate_segments, modulate_segments
|
||
|
|
|
||
|
|
|
||
|
|
class SegmentIndexContracts(unittest.TestCase):
|
||
|
|
def test_segment_index_covers_rows_and_reuses_layout(self):
|
||
|
|
try:
|
||
|
|
from h3_blackwell_runtime.h3_fusion import segment_index
|
||
|
|
except ImportError:
|
||
|
|
self.skipTest("Triton is not installed")
|
||
|
|
segments = [(0, 2, 1), (2, 5, 4), (5, 8, 0)]
|
||
|
|
first = segment_index(8, segments, torch.device("cpu"))
|
||
|
|
second = segment_index(8, segments, torch.device("cpu"))
|
||
|
|
self.assertIs(first, second)
|
||
|
|
torch.testing.assert_close(first, torch.tensor([1, 1, 4, 4, 4, 0, 0, 0], dtype=torch.int32))
|
||
|
|
|
||
|
|
def test_segment_index_rejects_gaps(self):
|
||
|
|
try:
|
||
|
|
from h3_blackwell_runtime.h3_fusion import segment_index
|
||
|
|
except ImportError:
|
||
|
|
self.skipTest("Triton is not installed")
|
||
|
|
with self.assertRaisesRegex(ValueError, "ordered, contiguous"):
|
||
|
|
segment_index(4, [(0, 2, 0), (3, 4, 1)], torch.device("cpu"))
|
||
|
|
|
||
|
|
|
||
|
|
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
|
||
|
|
class FusedElementwiseParity(unittest.TestCase):
|
||
|
|
@classmethod
|
||
|
|
def setUpClass(cls):
|
||
|
|
try:
|
||
|
|
from h3_blackwell_runtime.h3_fusion import (
|
||
|
|
fused_gate_add_, fused_modulate_, segment_index,
|
||
|
|
)
|
||
|
|
except ImportError as error:
|
||
|
|
raise unittest.SkipTest("Triton is not installed") from error
|
||
|
|
cls.fused_gate_add = staticmethod(fused_gate_add_)
|
||
|
|
cls.fused_modulate = staticmethod(fused_modulate_)
|
||
|
|
cls.segment_index = staticmethod(segment_index)
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
torch.manual_seed(440420)
|
||
|
|
self.segments = [(0, 19, 1), (19, 100, 4), (100, 201, 8), (201, 259, 2)]
|
||
|
|
self.row_index = self.segment_index(259, self.segments, torch.device("cuda"))
|
||
|
|
|
||
|
|
def test_modulation_is_bit_exact(self):
|
||
|
|
x = torch.randn(259, 1024, device="cuda", dtype=torch.bfloat16)
|
||
|
|
shift = torch.randn(9, 1024, device="cuda", dtype=torch.float32)
|
||
|
|
scale = torch.randn(9, 1024, device="cuda", dtype=torch.float32)
|
||
|
|
expected = modulate_segments(x, shift, scale, self.segments)
|
||
|
|
actual = self.fused_modulate(x.clone(), shift, scale, self.row_index)
|
||
|
|
self.assertTrue(torch.equal(actual, expected))
|
||
|
|
|
||
|
|
def test_gate_add_is_bit_exact(self):
|
||
|
|
residual = torch.randn(259, 1024, device="cuda", dtype=torch.bfloat16)
|
||
|
|
update = torch.randn_like(residual)
|
||
|
|
gate = torch.randn(9, 1024, device="cuda", dtype=torch.float32)
|
||
|
|
expected = gate_segments(residual, update, gate, self.segments)
|
||
|
|
actual = self.fused_gate_add(residual.clone(), update, gate, self.row_index)
|
||
|
|
self.assertTrue(torch.equal(actual, expected))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|