@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 199 lines 5.71 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """GLM-5.2 decode GEMM selection for unquantized BF16 on SM103."""
4
5 from __future__ import annotations
6
7 from dataclasses import dataclass
8 from typing import Literal
9
10 import torch
11 from torch import nn
12
13 import vllm.envs as envs
14 from vllm import _custom_ops as ops
15 from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
16 SkinnyGemmConfig,
17 shape_dynamic_skinny_gemm,
18 )
19 from vllm.model_executor.layers.linear import (
20 LinearBase,
21 UnquantizedLinearMethod,
22 )
23 from vllm.platforms import current_platform
24
25 Backend = Literal["cute", "dsv3_fused_a"]
26 ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
27
28
29 @dataclass(frozen=True, slots=True)
30 class GLM52ProjectionSpec:
31 n: int
32 k: int
33 cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...]
34 dsv3_tokens: frozenset[int] = frozenset()
35
36 def build_plan(self) -> dict[int, ResolvedCall]:
37 plan: dict[int, ResolvedCall] = {
38 num_tokens: ("cute", config) for num_tokens, config in self.cute_configs
39 }
40 plan.update(
41 (num_tokens, ("dsv3_fused_a", None)) for num_tokens in self.dsv3_tokens
42 )
43 return plan
44
45
46 GLM52_QKV_A_PROJECTION = GLM52ProjectionSpec(
47 n=2624,
48 k=6144,
49 cute_configs=(
50 (1, SkinnyGemmConfig(1, 128, 4, static_k=6144)),
51 (2, SkinnyGemmConfig(2, 128, 2)),
52 ),
53 dsv3_tokens=frozenset(range(3, 17)),
54 )
55
56 GLM52_Q_B_PROJECTION = GLM52ProjectionSpec(
57 n=2048,
58 k=2048,
59 cute_configs=(
60 (1, SkinnyGemmConfig(1, 128, 4, static_k=2048)),
61 (2, SkinnyGemmConfig(2, 64, 2, k_unroll=2)),
62 ),
63 dsv3_tokens=frozenset(range(3, 17)),
64 )
65
66 # The MTP eh_proj is a plain nn.Linear, so it gets its plan through
67 # build_glm52_plan rather than a quant_method swap. cuBLAS wins from M=4.
68 GLM52_EH_PROJECTION = GLM52ProjectionSpec(
69 n=6144,
70 k=12288,
71 cute_configs=(
72 (1, SkinnyGemmConfig(1, 256, 2, vector_width=4, static_k=12288)),
73 (2, SkinnyGemmConfig(2, 64, 2)),
74 (3, SkinnyGemmConfig(3, 64, 2)),
75 ),
76 )
77
78 GLM52_PROJECTIONS = {
79 (spec.n, spec.k): spec
80 for spec in (
81 GLM52_QKV_A_PROJECTION,
82 GLM52_Q_B_PROJECTION,
83 GLM52_EH_PROJECTION,
84 )
85 }
86
87
88 def _is_sm103() -> bool:
89 return current_platform.is_device_capability((10, 3))
90
91
92 def _is_supported_row_major(tensor: torch.Tensor) -> bool:
93 return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
94
95
96 def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
97 return (
98 not envs.VLLM_BATCH_INVARIANT
99 and _is_supported_row_major(x)
100 and _is_supported_row_major(weight)
101 and x.dtype == torch.bfloat16
102 and weight.dtype == torch.bfloat16
103 and x.is_cuda
104 and weight.is_cuda
105 and x.device == weight.device
106 and x.shape[1] == weight.shape[1]
107 )
108
109
110 def run_glm52_plan(
111 plan: dict[int, ResolvedCall] | None,
112 x: torch.Tensor,
113 weight: torch.Tensor,
114 ) -> torch.Tensor | None:
115 if plan is None or not _runtime_ok(x, weight):
116 return None
117 entry = plan.get(x.shape[0])
118 if entry is None:
119 return None
120
121 backend, config = entry
122 if backend == "cute":
123 if not shape_dynamic_skinny_gemm.is_available():
124 return None
125 return shape_dynamic_skinny_gemm(x, weight, config)
126
127 if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
128 return None
129 output = torch.empty(
130 (x.shape[0], weight.shape[0]),
131 dtype=x.dtype,
132 device=x.device,
133 )
134 ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True)
135 return output
136
137
138 def _request_warmup(dtype: torch.dtype, configs: set[SkinnyGemmConfig]) -> None:
139 if configs and shape_dynamic_skinny_gemm.is_available():
140 shape_dynamic_skinny_gemm.request_warmup_configs(dtype, configs)
141
142
143 def build_glm52_plan(
144 weight: torch.Tensor | None, dtype: torch.dtype
145 ) -> dict[int, ResolvedCall] | None:
146 """Plan for a weight the walk below cannot reach (a plain ``nn.Linear``)."""
147 if dtype != torch.bfloat16 or not _is_sm103():
148 return None
149 if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16:
150 return None
151 spec = GLM52_PROJECTIONS.get(tuple(weight.shape))
152 if spec is None:
153 return None
154 _request_warmup(dtype, {config for _, config in spec.cute_configs})
155 return spec.build_plan()
156
157
158 class GLM52LowLatencyLinearMethod(UnquantizedLinearMethod):
159 def __init__(self, plan: dict[int, ResolvedCall]) -> None:
160 super().__init__()
161 self._plan = plan
162
163 def apply(
164 self,
165 layer: nn.Module,
166 x: torch.Tensor,
167 bias: torch.Tensor | None = None,
168 ) -> torch.Tensor:
169 if bias is None:
170 output = run_glm52_plan(self._plan, x, layer.weight)
171 if output is not None:
172 return output
173 return super().apply(layer, x, bias)
174
175
176 def enable_glm52_low_latency_gemm(
177 module: nn.Module,
178 dtype: torch.dtype,
179 ) -> None:
180 if dtype != torch.bfloat16 or not _is_sm103():
181 return
182
183 warmup_configs: set[SkinnyGemmConfig] = set()
184 for child in module.modules():
185 if (
186 not isinstance(child, LinearBase)
187 or type(child.quant_method) is not UnquantizedLinearMethod
188 ):
189 continue
190 weight = getattr(child, "weight", None)
191 if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16:
192 continue
193 spec = GLM52_PROJECTIONS.get(tuple(weight.shape))
194 if spec is None:
195 continue
196 child.quant_method = GLM52LowLatencyLinearMethod(spec.build_plan())
197 warmup_configs.update(config for _, config in spec.cute_configs)
198
199 _request_warmup(dtype, warmup_configs)