For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
SwiGLUKernelConstants
struct SwiGLUKernelConstants[a_type: DType, b_type: DType, c_type: DType, transpose_b: Bool, config: FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]]
Compile-time constants for TMA descriptor creation and kernel launch.
Parameters
- a_type (
DType): Element data type of the A operand. - b_type (
DType): Element data type of the B operand. - c_type (
DType): Element data type of the C output. - transpose_b (
Bool): Whether B is transposed (always True for SwiGLU). - config (
FusedSwiGLUMatmulConfig[a_type, b_type, c_type, transpose_b]): Fused SwiGLU matmul config carrying tile shapes and pipeline stages.
Implemented traits
comptime members
a_tile_dim0
comptime a_tile_dim0 = (config.block_tile_shape[Int(0)] // config.cluster_shape[Int(1)])
AB_swapped
comptime AB_swapped = config.AB_swapped
ADescLayout
comptime ADescLayout = Layout[*?, *?]
ATileLayout
comptime ATileLayout = Layout[*?, *?]
b_tile_dim0
comptime b_tile_dim0 = (config.block_tile_shape[Int(1)] // (config.cluster_shape[Int(0)] // config))
BDescLayout
comptime BDescLayout = Layout[*?, *?]
Bias1DTile
comptime Bias1DTile = TileTensor[c_type, Layout[*?, *?], ImmutAnyOrigin]
Bias1DTileLayout
comptime Bias1DTileLayout = row_major[Int(1), config.mma_shape[Int(1)]]()
BK
comptime BK = config.block_tile_shape[Int(2)]
BM
comptime BM = config.block_tile_shape[Int(0)]
BN
comptime BN = config.block_tile_shape[Int(1)]
BTileLayout
comptime BTileLayout = Layout[*?, *?]
c_out_pair_tile_elems
comptime c_out_pair_tile_elems = (config.output_tile_shape[Int(0)] * (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)))
c_out_smem_inner
comptime c_out_smem_inner = (config.block_tile_shape[Int(0)] // Int(2)) if SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].AB_swapped else SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].HalfN
c_out_smem_offset
comptime c_out_smem_offset = (Int((add size_of[_SwiGLUSmem[a_type, b_type, c_type, transpose_b, config=config]](), 127)) & Int(-128))
c_out_tile_bytes
comptime c_out_tile_bytes = (Int((mul config.output_tile_shape[Int(0)], (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)), 2)) if config.AB_swapped and (eq config.cta_group, 2) and (eq config.mma_shape[Int(0)], 128) else Int((mul config.output_tile_shape[Int(0)], (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)))) * size_of[c_type]())
c_out_tile_elems
comptime c_out_tile_elems = (Int(2) * Int((mul config.output_tile_shape[Int(0)], (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2))))) if SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].AB_swapped and (config == Int(2)) and (config.mma_shape[Int(0)] == Int(128)) else SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].c_out_pair_tile_elems
c_store_m
comptime c_store_m = SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].OutputM if (config.mma_shape[Int(0)] == Int(256)) if (config.mma_shape[Int(0)] == Int(256)) else (config == Int(1)) or SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].AB_swapped else SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].BM
c_tma_inner
comptime c_tma_inner = max((config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)), Int(8)) if config.register_swiglu else (config.block_tile_shape[Int(0)] // Int(2)) if SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].AB_swapped else SwiGLUKernelConstants[a_type, b_type, c_type, transpose_b, config].HalfN
CDescLayout
comptime CDescLayout = Layout[*?, *?]
CLUSTER_M
comptime CLUSTER_M = config.cluster_shape[Int(0)]
CLUSTER_N
comptime CLUSTER_N = config.cluster_shape[Int(1)]
cta_group
comptime cta_group = config.cta_group
CTileLayout
comptime CTileLayout = Layout[*?, *?]
EPILOGUE_LOAD_THREADS
comptime EPILOGUE_LOAD_THREADS = WARP_SIZE if config.use_bias else Int(0)
HalfN
comptime HalfN = (config.output_tile_shape[Int(1)] // Int(2))
MMA_M
comptime MMA_M = config.mma_shape[Int(0)]
MMA_N
comptime MMA_N = config.mma_shape[Int(1)]
NUM_THREADS
comptime NUM_THREADS = (Int(224) + _resolve_warp_size() if config.use_bias else Int(0))
OutputM
comptime OutputM = config.output_tile_shape[Int(0)]
OutputN
comptime OutputN = config.output_tile_shape[Int(1)]
SmemType
comptime SmemType = _SwiGLUSmem[a_type, b_type, c_type, transpose_b, config=config]
total_smem_bytes
comptime total_smem_bytes = size_of[_SwiGLUSmem[a_type, b_type, c_type, transpose_b, config=config]]() if config.register_swiglu else (Int((add size_of[_SwiGLUSmem[a_type, b_type, c_type, transpose_b, config=config]](), 127) and -128) + Int((mul size_of[c_type](), Int((mul config.output_tile_shape[Int(0)], (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)), 2)) if config.AB_swapped and (eq config.cta_group, 2) and (eq config.mma_shape[Int(0)], 128) else Int((mul config.output_tile_shape[Int(0)], (config.block_tile_shape[Int(0)] // Int(2)) if config.AB_swapped else (config.output_tile_shape[Int(1)] // Int(2)))), 2)))
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!