prewarm5.py 727 B

123456789101112
  1. import os, time
  2. for k in ("FLASHINFER_JIT_VERBOSE", "FLASHINFER_JIT_DEBUG", "FLASHINFER_JIT_LINEINFO"):
  3. os.environ.pop(k, None) # these switch nvcc to debug/lineinfo flags -> cache mismatch vs the launcher
  4. t = time.time()
  5. from flashinfer.mla._sparse_mla_sm120 import get_sparse_mla_sm120_module as g
  6. try:
  7. g(); print("SPARSE-BUILT+LOADED %.1fs" % (time.time() - t), flush=True)
  8. except Exception as e:
  9. print("sparse build done, load raised (expected without GPU):", type(e).__name__, str(e)[:160], flush=True)
  10. from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8 as g1
  11. ic = getattr(g1(), "is_compiled"); ic = ic() if callable(ic) else ic
  12. print("mxfp8 is_compiled under runtime env:", ic, flush=True)