Module gemv_tc

Module gemv_tc 

Source
Expand description

2026-09-25: Routing of the narrow NVFP4 batched GEMV to its tensor-core kernels (w4a16_gemv_tc8 / w4a16_gemv_tc16 in kernels/gb10/common/w4a16_gemv_tc.cu).

Owner: model-layers ops. Invariants:

  • The tensor-core entries take the same arguments as the CUDA-core w4a16_gemv_batch* kernels, with their own grid (ceil(N / cols_per_cta)) and a 256-thread block.
  • A launch routes only when K % 128 == 0 and 1 <= m <= 16; otherwise the caller keeps its CUDA-core kernel (tc_route).
  • super::w4a16_gemv_batchm and tc_fixed_m both route through tc_kernel.
  • METRALE_NO_W4A16_TC set to any non-empty value, read once per process, turns the tensor-core path off (tc_enabled). The bench crate’s PERF_CONTROLS discloses it, and record_env_tests reads tc_enabled and wide_rows_enabled from this file.

Enums§

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

Constants§

NARROW_MAX_ROWS
2026-09-25: Row edges of the narrow-GEMV dispatch arms: NARROW_MAX_ROWS by default, and WIDE_MAX_ROWS (the tc16 reach) when wide_rows_enabled. Read through narrow_gemv_max_rows.
TC8_COLS_PER_CTA
2026-09-25: Columns per CTA, 8 * NT: tc8 has NT = 1, tc16 has NT = 2.
TC8_MAX_M
2026-09-25: Rows the w4a16_gemv_tc8 entry covers (its MT is 8).
TC16_COLS_PER_CTA
TC16_MAX_M
2026-09-25: Rows the w4a16_gemv_tc16 entry covers (its MT is 16).
TC_BLOCK
2026-09-25: Threads per CTA, TC_WARPS * 32 with TC_WARPS = 8 in the .cu; the warps split K.
WIDE_MAX_ROWS

Functions§

narrow_gemv_max_rows
2026-09-25: Row edge of the narrow-GEMV dispatch arms (see wide_rows_enabled).
tc_enabled
2026-09-25: True when METRALE_NO_W4A16_TC is unset or empty. Any non-empty value, 0 included, turns the tensor-core path off. Read once per process, so every launch, including one captured in a CUDA graph, sees the same choice.
tc_fixed_m
2026-09-25: The tensor-core route of the fixed-M launchers, taken only when wide_rows_enabled: Ok(true) when it launched, Ok(false) to keep the CUDA-core kernel.
tc_kernel
2026-09-25: The tensor-core kernel and grid x for this launch, or None to keep the CUDA-core kernel. It does not consult wide_rows_enabled; the callers’ row edges do.
tc_route
2026-09-25: The routing decision, without I/O. None keeps the caller’s CUDA-core kernel.
wide_rows_enabled
2026-09-25: The opt-in widening: true when METRALE_W4A16_TC_WIDE is set to a non-empty value and tc_enabled. Read once per process. It raises narrow_gemv_max_rows to 16 and lets tc_fixed_m route the fixed-M launchers (w4a16_gemv_batch2/3, w4a16_gemv_dual_batch2/3) to the tensor-core kernels.