Mike0021 commited on
Commit
e1a2a10
·
verified ·
1 Parent(s): 257cf81

diag: resolver introspection at load

Browse files
Files changed (1) hide show
  1. vsa_h3.py +11 -3
vsa_h3.py CHANGED
@@ -293,15 +293,23 @@ def _resolve_cuda_sparse_op():
293
  if _CUDA_SPARSE["resolved"]:
294
  return _CUDA_SPARSE["op"]
295
  _CUDA_SPARSE["resolved"] = True
 
 
 
 
 
296
  if os.environ.get("H3_VSA_CUDA", "1") != "1":
297
  return None
298
  try:
299
  from fastvideo_kernel import block_sparse_attn_sm100a as sm100a
300
 
301
- op = sm100a.block_sparse_attn_sm100a
302
- probe = sm100a._FWD_BY_BLOCK # noqa: F841 - import-time check that the extension loaded
303
  if not sm100a._HAS_VSA_SM100A:
304
- raise ImportError("fastvideo_kernel compiled without the sm100a VSA forward")
 
 
 
 
 
305
  _CUDA_SPARSE["op"] = op
306
  print("[vsa] CUDA sparse-attention kernel: fastvideo_kernel (sm100a fatbin on sm120)", flush=True)
307
  except Exception as error: # noqa: BLE001 - any failure means Triton
 
293
  if _CUDA_SPARSE["resolved"]:
294
  return _CUDA_SPARSE["op"]
295
  _CUDA_SPARSE["resolved"] = True
296
+ if os.environ.get("H3_VSA_CUDA", "1") != "1":
297
+ return None
298
+ if _CUDA_SPARSE["resolved"]:
299
+ return _CUDA_SPARSE["op"]
300
+ _CUDA_SPARSE["resolved"] = True
301
  if os.environ.get("H3_VSA_CUDA", "1") != "1":
302
  return None
303
  try:
304
  from fastvideo_kernel import block_sparse_attn_sm100a as sm100a
305
 
 
 
306
  if not sm100a._HAS_VSA_SM100A:
307
+ # Surface what the extension actually exposes so the logs show which pybind name this wheel has.
308
+ from fastvideo_kernel._C import fastvideo_kernel_ops as _ops
309
+
310
+ available = sorted(a for a in dir(_ops) if "sparse" in a.lower())
311
+ raise ImportError(f"sm100a forward flag unset; pybind exposes {available}")
312
+ op = sm100a.block_sparse_attn_sm100a
313
  _CUDA_SPARSE["op"] = op
314
  print("[vsa] CUDA sparse-attention kernel: fastvideo_kernel (sm100a fatbin on sm120)", flush=True)
315
  except Exception as error: # noqa: BLE001 - any failure means Triton