Module dense_gemv_tc

Module dense_gemv_tc 

Source
Expand description

2026-09-25: Routes the MTP drafter’s BF16 small-M GEMM rows to the tensor-core GEMV entries of dense_gemv_bf16_tc.cu.

Owner: model-layers ops. Invariants:

  • try_dense_gemv_tc launches only when METRALE_NO_MTP_TC is unset or empty, MIN_M <= m <= TC32_MAX_M, n > 0, k % K_STEP == 0 and a covering entry resolved; otherwise it returns Ok(false) and launches nothing.
  • METRALE_NO_MTP_TC is read once per process. Any non-empty value, 0 included, turns the path off.

The entries take dense_gemv_bf16_batchm’s arguments in the same order: A [M,K] contiguous, W [N,K], output rows out_stride elements apart. The warps of a CTA reduce through shared memory in a fixed order with no atomics, so the output is the same run to run.

Enums§

DtcKind
2026-09-25: Which tensor-core entry serves a launch.

Constants§

BLOCK
2026-09-25: Threads per CTA: DTC_WARPS (8) warps that split K.
K_STEP
2026-09-25: K per warp step (DTC_KB); a K that is not a multiple of it does not route.
MIN_M
2026-09-25: Narrowest M that routes; at M=1 the caller keeps its own kernel.
ROWS_PER_CTA
2026-09-25: Weight rows per CTA: 16 * NT, and every entry uses NT = 1.
TC8_MAX_M
2026-09-25: Widest M the tc8 entry covers: one 8-row token tile. The tc16 and tc32 entries run 2 and 4 tiles.
TC16_MAX_M
TC32_MAX_M

Functions§

kernel_for
2026-09-25: The tensor-core kernel and grid-x for this launch when the path is on and an entry routes, else None.
launch
2026-09-25: Launch the given entry with the given grid-x, without routing. try_dense_gemv_tc is the routed entry point.
mtp_tc_enabled
2026-09-25: Whether the path is on (METRALE_NO_MTP_TC unset or empty). Read once, so every launch in a process sees the same choice.
mtp_tc_from
2026-09-25: The METRALE_NO_MTP_TC rule over the looked-up value: the path is on unless the variable holds a non-empty value.
route
2026-09-25: Pure routing decision; None keeps the caller’s kernel.
try_dense_gemv_tc
2026-09-25: Run C[t] = A[t] @ W^T on a tensor-core entry if the path is on and the shape routes. Ok(false) means nothing was launched and the caller must run its own kernel.