biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix import path for tcgen05

Browse Source
This commit is contained in:
biondizzle
2026-05-23 06:34:30 +00:00
parent 44f8b4c5ce
commit 241aa0f334
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/unit/test_tmem_budget.py
View File

@@ -4,7 +4,8 @@ Prints find_tmem_tensor_col_offset(tOtO) and related shapes so we can
plan the SMEM-P path and verify TMEM fits in 512 columns at hd=512.
"""
import torch, math
import cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.nvgpu.tcgen05 as tcgen05
import cutlass, cutlass.cute as cute, cutlass.utils as utils
from cutlass.cute.nvgpu import tcgen05
from cutlass import BFloat16, Float32, LayoutEnum
from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset