Junda ChenandClaude Opus 5.5
6397822a5b
Exact-lane work from the FastH3 Spark agent team (2026-10-23): same computation as the shipped QAD-v2 Spark recipe, measured warm on GB10. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
36 lines
1.6 KiB
Python
36 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Private eager binding for the SM121 H3 SwiGLU/FP4 GEMM epilogue.
|
|
|
|
The Spark prebuild installs this module beside its fingerprinted native library.
|
|
Inputs are prequantized, with value/gate weight columns paired in memory.
|
|
"""
|
|
import ctypes
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
_library = ctypes.CDLL(str(Path(__file__).with_name('lib_fastvideo_h3_swiglu_gemm.so')))
|
|
_launch = _library.launch_gemm_swiglu
|
|
_launch.argtypes = ([ctypes.c_void_p] * 10 + [ctypes.c_int] * 3
|
|
+ [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int, ctypes.c_int, ctypes.c_void_p])
|
|
_launch.restype = ctypes.c_int
|
|
_workspaces: dict[torch.device, torch.Tensor] = {}
|
|
|
|
|
|
def run(a, b, a_scale, b_scale, alpha, global_scale, q, scales, amax=None, values=None,
|
|
raster=0, swizzle=1):
|
|
"""Write packed SwiGLU outputs; b and b_scale have contiguous transposes."""
|
|
workspace = _workspaces.get(a.device)
|
|
if workspace is None:
|
|
workspace = torch.empty(16 * 1024 * 1024, device=a.device, dtype=torch.uint8)
|
|
_workspaces[a.device] = workspace
|
|
rc = _launch(
|
|
a.data_ptr(), b.data_ptr(), a_scale.data_ptr(), b_scale.data_ptr(), alpha.data_ptr(),
|
|
global_scale.data_ptr(), q.data_ptr(), scales.data_ptr(),
|
|
0 if amax is None else amax.data_ptr(), 0 if values is None else values.data_ptr(),
|
|
a.shape[0], b.shape[1], a.shape[1] * 2, workspace.data_ptr(), workspace.numel(),
|
|
raster, swizzle, torch.cuda.current_stream(a.device).cuda_stream,
|
|
)
|
|
if rc:
|
|
raise RuntimeError(f'H3 SwiGLU GEMM returned CUDA/CUTLASS error {rc}')
|