fix: clarify import path setup for CuTeDSL

This commit is contained in:
2026-05-16 02:55:25 +00:00
parent f951d284e7
commit a3aa2d201e

View File

@@ -14,11 +14,15 @@ import os
import math
import torch
# Add CuTeDSL examples to path
# Add CuTeDSL examples to path — must be run from the examples directory
# so that 'from blackwell.kernel.moe.xxx' imports work
CUTLASS_ROOT = os.environ.get("CUTLASS_ROOT", "/root/cutlass")
CUTEDSL_EXAMPLES = os.path.join(CUTLASS_ROOT, "examples/python/CuTeDSL")
sys.path.insert(0, CUTEDSL_EXAMPLES)
# The imports below use 'from blackwell.kernel...' which resolves
# relative to CUTEDSL_EXAMPLES/cute/ (where the blackwell/ package lives)
import cutlass
import cutlass.cute as cute
import cutlass.torch as cutlass_torch