Module dense_gemv_tc 

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.