Skip to content

vllm.model_executor.warmup.jit_warmup_triton_helper

Classes:

Functions:

TritonWarmupTensor dataclass

Compile-only tensor metadata used by Triton warmup.

strides=None represents compact row-major storage. Pass explicit strides whenever the runtime tensor can be padded, transposed, or otherwise strided.

Source code in vllm/model_executor/warmup/jit_warmup_triton_helper.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@dataclass(frozen=True)
class TritonWarmupTensor:
    """Compile-only tensor metadata used by Triton warmup.

    ``strides=None`` represents compact row-major storage. Pass explicit strides
    whenever the runtime tensor can be padded, transposed, or otherwise strided.
    """

    dtype: Any
    aligned: bool = True
    shape: tuple[int, ...] = (1,)
    strides: tuple[int, ...] | None = None

    def data_ptr(self) -> int:
        return 0 if self.aligned else 1

    def ptr_range(self) -> int:
        return 0

    def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
        if self.strides is None:
            strides: list[int] = []
            stride = 1
            for size in reversed(self.shape):
                strides.append(stride)
                stride *= size
            result = tuple(reversed(strides))
        else:
            result = self.strides
        return result if dim is None else result[dim]

VllmTritonJitKernel

Bases: VllmJitKernel[CompileKeyT], Generic[CompileKeyT]

Triton owner whose runtime launch specification is reused for warmup.

Methods:

  • warmup_inputs –

    Return runtime-shaped inputs that reproduce one compile key.

Source code in vllm/model_executor/warmup/jit_warmup_triton_helper.py 
 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
class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
    """Triton owner whose runtime launch specification is reused for warmup."""

    kernel: ClassVar[Any]
    _warming = False

    @abstractmethod
    def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
        """Return runtime-shaped inputs that reproduce one compile key."""
        raise NotImplementedError

    def compile(self, compile_key: CompileKeyT) -> None:
        inputs = self.warmup_inputs(compile_key)
        self._warming = True
        try:
            cast(Callable[..., None], self)(**inputs)
        finally:
            self._warming = False

    @cached_property
    def _kernel_param_names(self) -> frozenset[str]:
        return frozenset(self.kernel.arg_names)

    def launch(
        self,
        grid: tuple[int, ...],
        inputs: Mapping[str, Any],
        /,
        **kwargs: Any,
    ) -> Any:
        for name, value in inputs.items():
            target = name if name in self._kernel_param_names else f"{name}_ptr"
            if target in self._kernel_param_names and target not in kwargs:
                kwargs[target] = value
        if self._warming:
            warmup = getattr(self.kernel, "warmup", None)
            assert warmup is not None
            return warmup(grid=(1,), **kwargs)
        return self.kernel[grid](**kwargs)

warmup_inputs(compile_key) abstractmethod

Return runtime-shaped inputs that reproduce one compile key.

Source code in vllm/model_executor/warmup/jit_warmup_triton_helper.py 
95
96
97
98
@abstractmethod
def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
    """Return runtime-shaped inputs that reproduce one compile key."""
    raise NotImplementedError

kernel_launcher(call_fn)

Launch a Triton kernel from a declarative __call__ specification.

Source code in vllm/model_executor/warmup/jit_warmup_triton_helper.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def kernel_launcher(
    call_fn: Callable[..., LaunchSpec],
) -> Callable[..., None]:
    """Launch a Triton kernel from a declarative ``__call__`` specification."""
    signature = inspect.signature(call_fn)

    @wraps(call_fn)
    def wrapper(
        self: VllmTritonJitKernel[Any],
        *args: Any,
        **kwargs: Any,
    ) -> None:
        grid, launch_kwargs = call_fn(self, *args, **kwargs)
        bound = signature.bind(self, *args, **kwargs)
        bound.apply_defaults()
        inputs = {
            name: value for name, value in bound.arguments.items() if name != "self"
        }
        self.launch(grid, inputs, **launch_kwargs)

    return wrapper

triton_scalar_specialization_rep(value)

Return an integer with the same default Triton JIT specialization.

For an ordinary integer argument, Triton's cache key contains its inferred type (i32, i64, or u64) and one of three value classes:

  • 1 is specialized as the exact constant 1.
  • Multiples of 16 receive a tt.divisibility = 16 attribute.
  • All other values have no value specialization.

Warmup only needs one concrete value for each cache-key class. This helper returns 1 for the exact-one class and otherwise returns a divisible or generic representative while preserving the inferred integer type.

This applies only to non-constexpr integer arguments using Triton's default specialization. Do not use it for arguments listed in do_not_specialize or do_not_specialize_on_alignment.

Source code in vllm/model_executor/warmup/jit_warmup_triton_helper.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def triton_scalar_specialization_rep(value: int) -> int:
    """Return an integer with the same default Triton JIT specialization.

    For an ordinary integer argument, Triton's cache key contains its inferred
    type (``i32``, ``i64``, or ``u64``) and one of three value classes:

    * ``1`` is specialized as the exact constant ``1``.
    * Multiples of 16 receive a ``tt.divisibility = 16`` attribute.
    * All other values have no value specialization.

    Warmup only needs one concrete value for each cache-key class. This helper
    returns ``1`` for the exact-one class and otherwise returns a divisible or
    generic representative while preserving the inferred integer type.

    This applies only to non-``constexpr`` integer arguments using Triton's
    default specialization. Do not use it for arguments listed in
    ``do_not_specialize`` or ``do_not_specialize_on_alignment``.
    """
    if value == 1:
        return 1

    if -(1 << 31) <= value < (1 << 31):
        divisible_rep = 16
        generic_rep = 2
    elif -(1 << 63) <= value < (1 << 63):
        divisible_rep = 1 << 31
        generic_rep = (1 << 31) + 1
    elif 0 <= value < (1 << 64):
        divisible_rep = 1 << 63
        generic_rep = (1 << 63) + 1
    else:
        raise OverflowError(f"Integer {value} is outside Triton's scalar range")

    return divisible_rep if value % 16 == 0 else generic_rep