From c4db339249c8ae7aa2adc9dbd144639e590df34f Mon Sep 17 00:00:00 2001 From: Will Lin Date: Sat, 22 Aug 2026 06:19:37 +0000 Subject: [PATCH] [perf] kernel: make sm100a sparse inference compile-safe --- .../block_sparse_attn_sm100a.py | 122 +++++++++++++++++- tests/test_block_sparse_sm100a.py | 69 ++++++++++ 2 files changed, 186 insertions(+), 5 deletions(-) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index 1ea406efc..a254b6586 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -84,6 +84,107 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: return True +@torch.library.custom_op( + "fastvideo_kernel::block_sparse_attn_sm100a_inference", + mutates_args=(), + device_types="cuda", +) +def _block_sparse_attn_sm100a_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> torch.Tensor: + """Opaque no-LSE launch used by the inference-only sm_100a route. + + The extension is exposed as a raw pybind function rather than a dispatcher + op. Calling it directly makes Dynamo descend through a Python/C++ boundary + that has no fake implementation, so ``torch.compile(fullgraph=True)`` + cannot capture a sparse H3 block. Keep that boundary inside this custom op; + its inputs have already been normalized by the public wrapper below. + """ + fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)] + sm_scale = 1.0 / (q.shape[-1]**0.5) + res = fwd(q, k, v, None, q2k_idx, q2k_num, variable_block_sizes, sm_scale, False) + return res[0] + + +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference") +def _block_sparse_attn_sm100a_inference_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> torch.Tensor: + # The C++ binding allocates its output with torch::empty_like(q). + return torch.empty_like(q) + + +@torch.library.custom_op( + "fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference", + mutates_args=(), + device_types="cuda", +) +def _block_sparse_attn_sm100a_from_mask_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> torch.Tensor: + """Opaque mask compaction plus no-LSE sm_100a launch for inference. + + H3 naturally produces a bool block map. Its Triton ``map_to_index`` call + must live behind the same opaque boundary as the raw pybind launch; + otherwise Dynamo sees that kernel before reaching the index-native custom + op and full-graph capture still fails. + """ + from fastvideo_kernel.triton_kernels.index import map_to_index + + q2k_idx, q2k_num = map_to_index(block_map) + return _block_sparse_attn_sm100a_inference( + q, + k, + v, + q2k_idx.to(torch.int32).contiguous(), + q2k_num.to(torch.int32).contiguous(), + variable_block_sizes, + ) + + +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference") +def _block_sparse_attn_sm100a_from_mask_inference_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> torch.Tensor: + return torch.empty_like(q) + + +def block_sparse_attn_sm100a_from_mask( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, None]: + """Inference forward from a bool block map, with compaction kept opaque.""" + out = _block_sparse_attn_sm100a_from_mask_inference( + q.contiguous(), + k.contiguous(), + v.contiguous(), + block_map.to(torch.bool).contiguous(), + variable_block_sizes.to(torch.int32).contiguous(), + ) + return out, None + + def block_sparse_attn_sm100a( q: torch.Tensor, k: torch.Tensor, @@ -92,13 +193,24 @@ def block_sparse_attn_sm100a( q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, need_lse: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor | None]: """Forward pass. Returns ``(out, lse)``; ``out`` has q's layout.""" - fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)] + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() idx = q2k_idx.to(torch.int32).contiguous() num = q2k_num.to(torch.int32).contiguous() vbs = variable_block_sizes.to(torch.int32).contiguous() + + if not need_lse: + # This is the production inference path. The custom op keeps the raw + # pybind launch opaque to Dynamo while its fake kernel carries output + # metadata through full-graph capture. + return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None + + # Preserve the established LSE-producing path for correctness tests and + # any future forward/backward pairing; only inference needs the opaque op. + fwd = _FWD_BY_BLOCK[_block_size(q, vbs)] sm_scale = 1.0 / (q.shape[-1]**0.5) - res = fwd(q.contiguous(), k.contiguous(), v.contiguous(), None, - idx, num, vbs, sm_scale, need_lse) - return (res[0], res[1]) if need_lse else (res[0], None) + res = fwd(q, k, v, None, idx, num, vbs, sm_scale, True) + return res[0], res[1] diff --git a/tests/test_block_sparse_sm100a.py b/tests/test_block_sparse_sm100a.py index 6f90f4017..479a59aff 100644 --- a/tests/test_block_sparse_sm100a.py +++ b/tests/test_block_sparse_sm100a.py @@ -49,6 +49,15 @@ def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed= return q, k, v, idx, num, vbs +def make_block_map(idx, num, batch, heads, num_blocks): + """Expand the test's compact rows to the bool map H3 produces.""" + block_map = torch.zeros((batch, heads, num_blocks, num_blocks), dtype=torch.bool, device=idx.device) + rows = block_map.view(-1, num_blocks) + for row in range(rows.shape[0]): + rows[row, idx[row, :int(num[row])].long()] = True + return block_map + + def reference(q, k, v, idx, num, vbs, block): """Dense attention restricted to the selected blocks, with padded keys masked.""" if not vsa.BHSD: @@ -123,6 +132,66 @@ def test_lse_is_not_vacuous(block): assert (got_lse.float() - (ref_lse + 1.0)).abs().max().item() > 0.5 +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_inference_no_lse_matches_eager_under_fullgraph_compile(block): + """The production no-LSE route is one opaque, fake-backed graph node.""" + q, k, v, idx, num, vbs = make_case(block, ragged=True) + + def inference_path(q, k, v, idx, num, vbs): + out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + assert lse is None + return out + + with torch.inference_mode(): + expected = inference_path(q, k, v, idx, num, vbs) + compiled = torch.compile(inference_path, backend="eager", fullgraph=True) + actual = compiled(q, k, v, idx, num, vbs) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_inference_custom_op_fake_matches_real(block): + """Pin the custom op schema and fake output metadata used by Dynamo.""" + q, k, v, idx, num, vbs = make_case(block, ragged=True) + torch.library.opcheck( + torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_inference.default, + (q, k, v, idx, num, vbs), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_mask_inference_matches_index_route_under_fullgraph_compile(block): + """H3's bool mask compaction and sm_100a launch stay in one opaque op.""" + batch, heads, num_blocks = 1, 4, 8 + q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) + block_map = make_block_map(idx, num, batch, heads, num_blocks) + + def inference_path(q, k, v, block_map, vbs): + out, lse = vsa.block_sparse_attn_sm100a_from_mask(q, k, v, block_map, vbs) + assert lse is None + return out + + with torch.inference_mode(): + expected, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + compiled = torch.compile(inference_path, backend="eager", fullgraph=True) + actual = compiled(q, k, v, block_map, vbs) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_mask_inference_custom_op_fake_matches_real(block): + """Pin schema/fake metadata for H3's mask-taking opaque boundary.""" + batch, heads, num_blocks = 1, 4, 8 + q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) + block_map = make_block_map(idx, num, batch, heads, num_blocks) + torch.library.opcheck( + torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_from_mask_inference.default, + (q, k, v, block_map, vbs), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) + + def test_unsupported_is_rejected(): q, _, _, _, _, vbs = make_case(64) assert not vsa.is_supported(q.float(), vbs) # wrong dtype