Skip to content

vllm.kernels.helion.ops.silu_and_mul_per_block_quant

Functions:

  • pick_config –

    Pick the best pre-tuned config for the given input shape.

pick_config(args, config_keys)

Pick the best pre-tuned config for the given input shape.

Selection strategy
  1. Find the closest intermediate_size among available configs (exact match preferred).
  2. Find the closest group_size among available configs (exact match preferred).
  3. Among the num_tokens values tuned for that intermediate_size and group_size, pick the smallest num_tokens >= the input's num_tokens. If the input is larger than all available num_tokens, fall back to the largest.
Source code in vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
    """Pick the best pre-tuned config for the given input shape.

    Selection strategy:
      1. Find the closest intermediate_size among available configs
         (exact match preferred).
      2. Find the closest group_size among available configs
         (exact match preferred).
      3. Among the num_tokens values tuned for that intermediate_size and group_size,
         pick the smallest num_tokens >= the input's num_tokens. If the input is
         larger than all available num_tokens, fall back to the largest.
    """

    if not config_keys:
        return None

    result, _, _, group_size, *_ = args
    num_tokens, intermediate_size = result.shape

    cache_key = (num_tokens, group_size, intermediate_size)
    cached = _pick_cache.get(cache_key)
    if cached is not None:
        return cached

    configs: dict[int, dict[int, list[int]]] = {}
    for key in config_keys:
        if key.is_default():
            continue
        configs.setdefault(key["intermediate_size"], {}).setdefault(
            key["group_size"], []
        ).append(key["num_tokens"])

    if not configs:
        return None

    best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size))
    best_group_size = min(
        configs[best_intermediate_size], key=lambda s: abs(s - group_size)
    )
    available_num_tokens = sorted(configs[best_intermediate_size][best_group_size])
    best_num_tokens = next(
        (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
    )

    result = CaseKey(
        {
            "intermediate_size": best_intermediate_size,
            "group_size": best_group_size,
            "num_tokens": best_num_tokens,
        }
    )
    _pick_cache[cache_key] = result
    return result