IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

gemv_gpu_dispatch

def gemv_gpu_dispatch[transpose_b: Bool = False, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, pdl_level: PDLLevel = PDLLevel.ON, tile_m: Int = Int(1)](kernel_func: GEMVAlgorithm, c: TileTensor[Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], ctx: DeviceContext)

Launches the GPU GEMV kernel indicated by kernel_func with appropriate grid and block dims.

Translates a GEMVAlgorithm variant into a concrete kernel call with shape-derived launch parameters, handling input/output layout transformation where needed.

Parameters:

Args:

def gemv_gpu_dispatch[EpilogueFnType: ElementwiseEpilogueFn, //, transpose_b: Bool = False, pdl_level: PDLLevel = PDLLevel.ON, tile_m: Int = Int(1)](kernel_func: GEMVAlgorithm, c: TileTensor[Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], epilogue_fn: EpilogueFnType, ctx: DeviceContext)

Launches the GEMV kernel indicated by kernel_func, storing through epilogue_fn.

Parameters:

  • ​EpilogueFnType (ElementwiseEpilogueFn): Type of epilogue_fn (inferred).
  • ​transpose_b (Bool): When True, B is treated as transposed (N, K) row-major.
  • ​pdl_level (PDLLevel): Programmatic dependent launch level.
  • ​tile_m (Int): Number of output rows processed per CTA (used by GEMV_SPLIT_K).

Args:

Was this page helpful?