fix: guard tma_p creation with gP check
This commit is contained in:
@@ -115,16 +115,16 @@ class FmhaKernel:
|
||||
tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
|
||||
|
||||
# SMEM-P: gP buffer and TMA for P (GMEM→SMEM via TMA)
|
||||
# gP is created by the caller and passed as a parameter
|
||||
# tma_p is created from the traced gP tensor
|
||||
if self.use_smem_p and gP is not None:
|
||||
p_s = cute.slice_(self.p_smem_s,(None,None,None,0))
|
||||
tma_p,gP = cute.nvgpu.make_tiled_tma_atom_A(
|
||||
utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, pv_mma.thr_id),
|
||||
gP, p_s, self.qk_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape
|
||||
)
|
||||
elif not self.use_smem_p:
|
||||
tma_p = tma_q # dummy, dead code
|
||||
else:
|
||||
tma_p = tma_q # dummy
|
||||
raise ValueError("use_smem_p=True but no gP provided")
|
||||
# Always create a valid mLSE tensor for the kernel.
|
||||
# CuTeDSL doesn't support None parameters in @cute.kernel.
|
||||
# For normalize=True, mLSE is unused (dead-code-eliminated by compiler).
|
||||
|
||||
Reference in New Issue
Block a user