Skip to content

vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm_plans

Measured skinny GEMM plans for unquantized BF16 linear layers.

Plans are keyed by device capability, local (N, K) weight shape and token count. A layer whose shape is listed runs the planned config for that token count; other token counts use F.linear. Entries are kept only where the skinny GEMM beats cuBLAS by more than 5% under CUDA graph replay.

Functions:

get_skinny_gemm_plan(n, k)

Return the token-count plan for an (N, K) weight on this device.

Source code in vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm_plans.py 
71
72
73
74
75
76
77
78
def get_skinny_gemm_plan(n: int, k: int) -> SkinnyGemmPlan | None:
    """Return the token-count plan for an ``(N, K)`` weight on this device."""
    plans = SKINNY_GEMM_PLANS_BY_CAPABILITY.get(
        current_platform.get_device_capability()
    )
    if plans is None or not shape_dynamic_skinny_gemm.is_available():
        return None
    return plans.get((n, k))