Merged
feat(quant): add per_token_cast torch.compile support #53
daisysjAdd on 13 days ago
feat(quant): add per_token_cast torch.compile support #53
Merged
Merged fromdaisysj/ops-tilelanggraph_adaptintocann/ops-tilelangmaster
Total: 5 files changed+897-18
| @@ -1,13 +1,19 @@ | |||
| 1 | +from typing import NamedTuple | ||
| 2 | + | ||
| 1 | import torch | 3 | import torch |
| 2 | 4 | ||
| 3 | -from cann_ops_tilelang.config import require_ascend, should_print_kernel_source | 5 | +from cann_ops_tilelang.config import get_amax_clamp_for_quant, require_ascend, should_print_kernel_source |
| 4 | 6 | ||
| 5 | from .common import ( | 7 | from .common import ( |
| 8 | + CastInputConfig, | ||
| 9 | + CastOutputConfig, | ||
| 6 | alloc_scaling_factors, | 10 | alloc_scaling_factors, |
| 7 | cast_epilogue, | 11 | cast_epilogue, |
| 8 | get_cast_input_and_config, | 12 | get_cast_input_and_config, |
| 9 | get_cast_output_config, | 13 | get_cast_output_config, |
| 14 | + get_cast_output_config_impl, | ||
| 10 | get_logical_hidden, | 15 | get_logical_hidden, |
| 16 | + get_packed_ue8m0_torch_dtype, | ||
| 11 | get_physical_hidden, | 17 | get_physical_hidden, |
| 12 | get_sf_shape, | 18 | get_sf_shape, |
| 13 | ) | 19 | ) |
| @@ -15,6 +21,53 @@ from .per_token_cast_kernel import get_per_token_cast_kernel_asc, tall_prequant_ | |||
| 15 | from .types import QuantTensor | 21 | from .types import QuantTensor |
| 16 | 22 | ||
| 17 | 23 | ||
| 24 | +class _PreparedCast(NamedTuple): | ||
| 25 | + x: torch.Tensor | ||
| 26 | + x_sf: torch.Tensor | None | ||
| 27 | + in_config: CastInputConfig | ||
| 28 | + out_config: CastOutputConfig | ||
| 29 | + num_tokens: int | ||
| 30 | + hidden: int | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def _validate_per_token_cast( | ||
| 34 | + x: torch.Tensor | QuantTensor, | ||
| 35 | + fmt: str, | ||
| 36 | + num_per_channels: int, | ||
| 37 | + x_block_size: tuple[int, int] | None, | ||
| 38 | + use_tma_aligned_col_major_sf: bool, | ||
| 39 | + round_sf: bool, | ||
| 40 | + use_packed_ue8m0: bool, | ||
| 41 | + clamp_min_value: float | None = None, | ||
| 42 | +) -> _PreparedCast: | ||
| 43 | + """Validate metadata without environment probes, JIT or device execution.""" | ||
| 44 | + assert fmt in ('e4m3', 'e2m1') | ||
| 45 | + | ||
| 46 | + x, x_sf, in_config = get_cast_input_and_config(x, x_block_size) | ||
| 47 | + if x.device.type != 'npu': | ||
| 48 | + raise ValueError(f'x must be an NPU tensor, got {x.device}.') | ||
| 49 | + if clamp_min_value is None: | ||
| 50 | + out_config = get_cast_output_config( | ||
| 51 | + fmt, (1, num_per_channels), use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0 | ||
| 52 | + ) | ||
| 53 | + else: | ||
| 54 | + out_config = get_cast_output_config_impl( | ||
| 55 | + fmt, | ||
| 56 | + (1, num_per_channels), | ||
| 57 | + use_tma_aligned_col_major_sf, | ||
| 58 | + round_sf, | ||
| 59 | + use_packed_ue8m0, | ||
| 60 | + clamp_min_value, | ||
| 61 | + ) | ||
| 62 | + | ||
| 63 | + assert x.dim() == 2 | ||
| 64 | + assert x.stride(1) == 1 | ||
| 65 | + num_tokens, hidden = x.shape | ||
| 66 | + hidden = get_logical_hidden(hidden, x.dtype) | ||
| 67 | + assert num_per_channels in (16, 32, 64, 128) or (num_per_channels == hidden and hidden % 64 == 0) | ||
| 68 | + return _PreparedCast(x, x_sf, in_config, out_config, num_tokens, hidden) | ||
| 69 | + | ||
| 70 | + | ||
| 18 | def per_token_cast_impl( | 71 | def per_token_cast_impl( |
| 19 | x: torch.Tensor | QuantTensor, | 72 | x: torch.Tensor | QuantTensor, |
| 20 | fmt: str, | 73 | fmt: str, |
| @@ -28,22 +81,19 @@ def per_token_cast_impl( | |||
| 28 | sf: torch.Tensor | None = None, | 81 | sf: torch.Tensor | None = None, |
| 29 | out_x: torch.Tensor | None = None, | 82 | out_x: torch.Tensor | None = None, |
| 30 | out_sf: torch.Tensor | None = None, | 83 | out_sf: torch.Tensor | None = None, |
| 84 | + clamp_min_value: float | None = None, | ||
| 31 | ) -> torch.Tensor | QuantTensor: | 85 | ) -> torch.Tensor | QuantTensor: |
| 32 | require_ascend() | 86 | require_ascend() |
| 33 | - assert fmt in ('e4m3', 'e2m1') | 87 | + x, x_sf, in_config, out_config, num_tokens, hidden = _validate_per_token_cast( |
| 34 | - | 88 | + x, |
| 35 | - x, x_sf, in_config = get_cast_input_and_config(x, x_block_size) | 89 | + fmt, |
| 36 | - if x.device.type != 'npu': | 90 | + num_per_channels, |
| 37 | - raise ValueError(f'x must be an NPU tensor, got {x.device}.') | 91 | + x_block_size, |
| 38 | - out_config = get_cast_output_config( | 92 | + use_tma_aligned_col_major_sf, |
| 39 | - fmt, (1, num_per_channels), use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0 | 93 | + round_sf, |
| 94 | + use_packed_ue8m0, | ||
| 95 | + clamp_min_value, | ||
| 40 | ) | 96 | ) |
| 41 | - | ||
| 42 | - assert x.dim() == 2 | ||
| 43 | - assert x.stride(1) == 1 | ||
| 44 | - num_tokens, hidden = x.shape | ||
| 45 | - hidden = get_logical_hidden(hidden, x.dtype) | ||
| 46 | - assert num_per_channels in (16, 32, 64, 128) or (num_per_channels == hidden and hidden % 64 == 0) | ||
| 47 | assert sf is None or out_sf is None | 97 | assert sf is None or out_sf is None |
| 48 | 98 | ||
| 49 | # The kernel is built for a hidden size but not for a token count, so | 99 | # The kernel is built for a hidden size but not for a token count, so |
| @@ -88,15 +138,105 @@ def per_token_cast_impl( | |||
| 88 | if num_tokens > 0: | 138 | if num_tokens > 0: |
| 89 | kernel(x, x_sf, out_x, out_sf) | 139 | kernel(x, x_sf, out_x, out_sf) |
| 90 | 140 | ||
| 141 | + # Cast-only does not return SF. Avoid unnecessary dtype views of the | ||
| 142 | + # read-only caller buffer, especially empty inputs rebuilt by graph capture. | ||
| 143 | + if sf is not None and not sf_only: | ||
| 144 | + return out_x | ||
| 91 | out_sf = cast_epilogue(out_sf, out_config) | 145 | out_sf = cast_epilogue(out_sf, out_config) |
| 92 | 146 | ||
| 93 | if sf_only: | 147 | if sf_only: |
| 94 | return out_sf | 148 | return out_sf |
| 95 | - if sf is not None: | ||
| 96 | - return out_x | ||
| 97 | return out_x, out_sf | 149 | return out_x, out_sf |
| 98 | 150 | ||
| 99 | 151 | ||
| 152 | +def _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value, x_sf=None, x_block_size=None): | ||
| 153 | + """Decode graph-schema flags into the existing eager implementation's keywords. | ||
| 154 | + | ||
| 155 | + The bool list groups related static options, in order: TMA layout, scale | ||
| 156 | + rounding, UE8M0 packing, stochastic casting. Clamp stays an explicit input. | ||
| 157 | + """ | ||
| 158 | + tma, rounded, packed, stochastic = options | ||
| 159 | + return dict( | ||
| 160 | + x=x if x_sf is None else (x, x_sf), | ||
| 161 | + fmt=fmt, | ||
| 162 | + num_per_channels=num_per_channels, | ||
| 163 | + x_block_size=None if x_block_size is None else tuple(x_block_size), | ||
| 164 | + use_tma_aligned_col_major_sf=tma, | ||
| 165 | + round_sf=rounded, | ||
| 166 | + use_packed_ue8m0=packed, | ||
| 167 | + stochastic_cast=stochastic, | ||
| 168 | + clamp_min_value=clamp_min_value, | ||
| 169 | + ) | ||
| 170 | + | ||
| 171 | + | ||
| 172 | +def _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value, x_sf=None, x_block_size=None): | ||
| 173 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value, x_sf, x_block_size) | ||
| 174 | + del kwargs['stochastic_cast'] | ||
| 175 | + return _validate_per_token_cast(**kwargs) | ||
| 176 | + | ||
| 177 | + | ||
| 178 | +def _allocate_cast_data(prepared): | ||
| 179 | + return torch.empty( | ||
| 180 | + (prepared.num_tokens, get_physical_hidden(prepared.hidden, prepared.out_config.torch_dtype)), | ||
| 181 | + dtype=prepared.out_config.torch_dtype, | ||
| 182 | + device=prepared.x.device, | ||
| 183 | + ) | ||
| 184 | + | ||
| 185 | + | ||
| 186 | +def _allocate_cast_scale(prepared): | ||
| 187 | + scale = alloc_scaling_factors((prepared.num_tokens, prepared.hidden), prepared.out_config, prepared.x.device) | ||
| 188 | + return cast_epilogue(scale, prepared.out_config) | ||
| 189 | + | ||
| 190 | + | ||
| 191 | +def _validate_scale_output(prepared, out_sf): | ||
| 192 | + config = prepared.out_config | ||
| 193 | + assert not config.use_tma_aligned_col_major_sf, 'Not supported yet' | ||
| 194 | + raw_sf = out_sf.view(torch.uint8) if config.use_packed_ue8m0 else out_sf | ||
| 195 | + assert raw_sf.shape == get_sf_shape((prepared.num_tokens, prepared.hidden), config) | ||
| 196 | + | ||
| 197 | + | ||
| 198 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast', mutates_args=()) | ||
| 199 | +def _compiled_per_token_cast( | ||
| 200 | + x: torch.Tensor, | ||
| 201 | + x_sf: torch.Tensor | None, | ||
| 202 | + fmt: str, | ||
| 203 | + num_per_channels: int, | ||
| 204 | + x_block_size: list[int] | None, | ||
| 205 | + options: list[bool], | ||
| 206 | + clamp_min_value: float, | ||
| 207 | +) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 208 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value, x_sf, x_block_size) | ||
| 209 | + return per_token_cast_impl(**kwargs) | ||
| 210 | + | ||
| 211 | + | ||
| 212 | +@_compiled_per_token_cast.register_fake | ||
| 213 | +def _per_token_cast_fake(x, x_sf, fmt, num_per_channels, x_block_size, options, clamp_min_value): | ||
| 214 | + prepared = _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value, x_sf, x_block_size) | ||
| 215 | + return _allocate_cast_data(prepared), _allocate_cast_scale(prepared) | ||
| 216 | + | ||
| 217 | + | ||
| 218 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_out', mutates_args=('out_x', 'out_sf')) | ||
| 219 | +def _compiled_per_token_cast_out( | ||
| 220 | + x: torch.Tensor, | ||
| 221 | + x_sf: torch.Tensor | None, | ||
| 222 | + fmt: str, | ||
| 223 | + num_per_channels: int, | ||
| 224 | + x_block_size: list[int] | None, | ||
| 225 | + options: list[bool], | ||
| 226 | + clamp_min_value: float, | ||
| 227 | + out_x: torch.Tensor, | ||
| 228 | + out_sf: torch.Tensor, | ||
| 229 | +) -> None: | ||
| 230 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value, x_sf, x_block_size) | ||
| 231 | + per_token_cast_impl(**kwargs, out_x=out_x, out_sf=out_sf) | ||
| 232 | + | ||
| 233 | + | ||
| 234 | +@_compiled_per_token_cast_out.register_fake | ||
| 235 | +def _per_token_cast_out_fake(x, x_sf, fmt, num_per_channels, x_block_size, options, clamp_min_value, out_x, out_sf): | ||
| 236 | + prepared = _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value, x_sf, x_block_size) | ||
| 237 | + _validate_scale_output(prepared, out_sf) | ||
| 238 | + | ||
| 239 | + | ||
| 100 | def per_token_cast( | 240 | def per_token_cast( |
| 101 | x: torch.Tensor | QuantTensor, | 241 | x: torch.Tensor | QuantTensor, |
| 102 | fmt: str, | 242 | fmt: str, |
| @@ -123,7 +263,39 @@ def per_token_cast( | |||
| 123 | 263 | ||
| 124 | Returns: | 264 | Returns: |
| 125 | A tuple ``(out, out_sf)`` with quantized output and sf factors. | 265 | A tuple ``(out, out_sf)`` with quantized output and sf factors. |
| 266 | + | ||
| 267 | + The compile path uses custom operators with metadata-only fake implementations. | ||
| 268 | + The allocation and out-buffer paths have separate mutation contracts. Kernel | ||
| 269 | + math and stochastic rounding are unchanged. Use fixed-shape npugraph_ex | ||
| 270 | + execution in inference_mode; no autograd formula is registered. | ||
| 126 | """ | 271 | """ |
| 272 | + if torch.compiler.is_compiling(): | ||
| 273 | + # Flatten QuantTensor only at the internal registration boundary. | ||
| 274 | + x_data, x_sf = x if isinstance(x, tuple) else (x, None) | ||
| 275 | + block_size = None if x_block_size is None else list(x_block_size) | ||
| 276 | + # Explicit configuration input: changing the clamp must not be hidden | ||
| 277 | + # behind the opaque custom op or a cached configuration object. | ||
| 278 | + clamp_min_value = get_amax_clamp_for_quant(fmt) | ||
| 279 | + args = ( | ||
| 280 | + x_data, | ||
| 281 | + x_sf, | ||
| 282 | + fmt, | ||
| 283 | + num_per_channels, | ||
| 284 | + block_size, | ||
| 285 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, stochastic_cast], | ||
| 286 | + clamp_min_value, | ||
| 287 | + ) | ||
| 288 | + if out is None: | ||
| 289 | + return _compiled_per_token_cast(*args) | ||
| 290 | + out_x, out_sf = out | ||
| 291 | + _compiled_per_token_cast_out(*args, out_x, out_sf) | ||
| 292 | + # custom_op cannot return an input alias. Reconstruct the original | ||
| 293 | + # public return outside that boundary, preserving the packed-SF view. | ||
| 294 | + returned_sf = out_sf.view(torch.uint8) if use_packed_ue8m0 else out_sf | ||
| 295 | + if use_packed_ue8m0: | ||
| 296 | + returned_sf = returned_sf.view(get_packed_ue8m0_torch_dtype()) | ||
| 297 | + return out_x, returned_sf | ||
| 298 | + | ||
| 127 | out_x, out_sf = (None, None) if out is None else out | 299 | out_x, out_sf = (None, None) if out is None else out |
| 128 | return per_token_cast_impl( | 300 | return per_token_cast_impl( |
| 129 | x, | 301 | x, |
| @@ -139,6 +311,82 @@ def per_token_cast( | |||
| 139 | ) | 311 | ) |
| 140 | 312 | ||
| 141 | 313 | ||
| 314 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_with_sf_only', mutates_args=()) | ||
| 315 | +def _compiled_per_token_cast_with_sf_only( | ||
| 316 | + x: torch.Tensor, | ||
| 317 | + fmt: str, | ||
| 318 | + num_per_channels: int, | ||
| 319 | + options: list[bool], | ||
| 320 | + clamp_min_value: float, | ||
| 321 | +) -> torch.Tensor: | ||
| 322 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 323 | + return per_token_cast_impl(**kwargs, sf_only=True) | ||
| 324 | + | ||
| 325 | + | ||
| 326 | +@_compiled_per_token_cast_with_sf_only.register_fake | ||
| 327 | +def _per_token_cast_with_sf_only_fake(x, fmt, num_per_channels, options, clamp_min_value): | ||
| 328 | + prepared = _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 329 | + return _allocate_cast_scale(prepared) | ||
| 330 | + | ||
| 331 | + | ||
| 332 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_with_sf_only_out', mutates_args=('out_sf',)) | ||
| 333 | +def _compiled_per_token_cast_with_sf_only_out( | ||
| 334 | + x: torch.Tensor, | ||
| 335 | + fmt: str, | ||
| 336 | + num_per_channels: int, | ||
| 337 | + options: list[bool], | ||
| 338 | + clamp_min_value: float, | ||
| 339 | + out_sf: torch.Tensor, | ||
| 340 | +) -> None: | ||
| 341 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 342 | + per_token_cast_impl(**kwargs, sf_only=True, out_sf=out_sf) | ||
| 343 | + | ||
| 344 | + | ||
| 345 | +@_compiled_per_token_cast_with_sf_only_out.register_fake | ||
| 346 | +def _per_token_cast_with_sf_only_out_fake(x, fmt, num_per_channels, options, clamp_min_value, out_sf): | ||
| 347 | + prepared = _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 348 | + _validate_scale_output(prepared, out_sf) | ||
| 349 | + | ||
| 350 | + | ||
| 351 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_with_precomputed_sf', mutates_args=()) | ||
| 352 | +def _compiled_per_token_cast_with_precomputed_sf( | ||
| 353 | + x: torch.Tensor, | ||
| 354 | + sf: torch.Tensor, | ||
| 355 | + fmt: str, | ||
| 356 | + num_per_channels: int, | ||
| 357 | + options: list[bool], | ||
| 358 | + clamp_min_value: float, | ||
| 359 | +) -> torch.Tensor: | ||
| 360 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 361 | + return per_token_cast_impl(**kwargs, sf=sf) | ||
| 362 | + | ||
| 363 | + | ||
| 364 | +@_compiled_per_token_cast_with_precomputed_sf.register_fake | ||
| 365 | +def _per_token_cast_with_precomputed_sf_fake(x, sf, fmt, num_per_channels, options, clamp_min_value): | ||
| 366 | + prepared = _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 367 | + return _allocate_cast_data(prepared) | ||
| 368 | + | ||
| 369 | + | ||
| 370 | +@torch.library.custom_op('cann_ops_tilelang::per_token_cast_with_precomputed_sf_out', mutates_args=('out_x',)) | ||
| 371 | +def _compiled_per_token_cast_with_precomputed_sf_out( | ||
| 372 | + x: torch.Tensor, | ||
| 373 | + sf: torch.Tensor, | ||
| 374 | + fmt: str, | ||
| 375 | + num_per_channels: int, | ||
| 376 | + options: list[bool], | ||
| 377 | + clamp_min_value: float, | ||
| 378 | + out_x: torch.Tensor, | ||
| 379 | +) -> None: | ||
| 380 | + kwargs = _registered_cast_kwargs(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 381 | + per_token_cast_impl(**kwargs, sf=sf, out_x=out_x) | ||
| 382 | + | ||
| 383 | + | ||
| 384 | +@_compiled_per_token_cast_with_precomputed_sf_out.register_fake | ||
| 385 | +def _per_token_cast_with_precomputed_sf_out_fake(x, sf, fmt, num_per_channels, options, clamp_min_value, out_x): | ||
| 386 | + # Caller owns output storage; only validate input metadata here. | ||
| 387 | + _prepare_registered_cast(x, fmt, num_per_channels, options, clamp_min_value) | ||
| 388 | + | ||
| 389 | + | ||
| 142 | def per_token_cast_with_sf_only( | 390 | def per_token_cast_with_sf_only( |
| 143 | x: torch.Tensor, | 391 | x: torch.Tensor, |
| 144 | fmt: str, | 392 | fmt: str, |
| @@ -148,7 +396,26 @@ def per_token_cast_with_sf_only( | |||
| 148 | use_packed_ue8m0: bool = False, | 396 | use_packed_ue8m0: bool = False, |
| 149 | out_sf: torch.Tensor | None = None, | 397 | out_sf: torch.Tensor | None = None, |
| 150 | ) -> torch.Tensor: | 398 | ) -> torch.Tensor: |
| 151 | - """Return only the per-token scaling factors for an FP8/FP4 cast.""" | 399 | + """Return only the per-token scaling factors for an FP8/FP4 cast. |
| 400 | + | ||
| 401 | + Compile registration preserves allocation and caller-owned output modes. | ||
| 402 | + Fixed-shape inference graph execution is supported with npugraph_ex. | ||
| 403 | + """ | ||
| 404 | + if torch.compiler.is_compiling(): | ||
| 405 | + args = ( | ||
| 406 | + x, | ||
| 407 | + fmt, | ||
| 408 | + num_per_channels, | ||
| 409 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, False], | ||
| 410 | + get_amax_clamp_for_quant(fmt), | ||
| 411 | + ) | ||
| 412 | + if out_sf is None: | ||
| 413 | + return _compiled_per_token_cast_with_sf_only(*args) | ||
| 414 | + _compiled_per_token_cast_with_sf_only_out(*args, out_sf) | ||
| 415 | + # Return the same storage/view as the eager SF epilogue, outside the op. | ||
| 416 | + if use_packed_ue8m0: | ||
| 417 | + return out_sf.view(torch.uint8).view(get_packed_ue8m0_torch_dtype()) | ||
| 418 | + return out_sf | ||
| 152 | return per_token_cast_impl( | 419 | return per_token_cast_impl( |
| 153 | x, | 420 | x, |
| 154 | fmt, | 421 | fmt, |
| @@ -172,7 +439,26 @@ def per_token_cast_with_precomputed_sf( | |||
| 172 | use_packed_ue8m0: bool = False, | 439 | use_packed_ue8m0: bool = False, |
| 173 | out: torch.Tensor | None = None, | 440 | out: torch.Tensor | None = None, |
| 174 | ) -> torch.Tensor: | 441 | ) -> torch.Tensor: |
| 175 | - """Cast with caller-provided per-token scaling factors.""" | 442 | + """Cast with caller-provided per-token scaling factors. |
| 443 | + | ||
| 444 | + The supplied sf remains read-only; only a caller-provided out is mutated. | ||
| 445 | + sf uses the raw kernel layout: transpose a TMA-aligned public scale tensor | ||
| 446 | + first, then view packed UE8M0 storage as uint8 when packing is enabled. | ||
| 447 | + Fixed-shape inference graph execution is supported with npugraph_ex. | ||
| 448 | + """ | ||
| 449 | + if torch.compiler.is_compiling(): | ||
| 450 | + args = ( | ||
| 451 | + x, | ||
| 452 | + sf, | ||
| 453 | + fmt, | ||
| 454 | + num_per_channels, | ||
| 455 | + [use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, False], | ||
| 456 | + get_amax_clamp_for_quant(fmt), | ||
| 457 | + ) | ||
| 458 | + if out is None: | ||
| 459 | + return _compiled_per_token_cast_with_precomputed_sf(*args) | ||
| 460 | + _compiled_per_token_cast_with_precomputed_sf_out(*args, out) | ||
| 461 | + return out | ||
| 176 | return per_token_cast_impl( | 462 | return per_token_cast_impl( |
| 177 | x, | 463 | x, |
| 178 | fmt, | 464 | fmt, |
| @@ -532,3 +532,44 @@ def test_load_kv_cache_llama_benchmark( | |||
| 532 | time_us=t_us, | 532 | time_us=t_us, |
| 533 | extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)}, | 533 | extras={'bandwidth_gbs': total_bytes / 1000000000.0 / (t_us / 1000000.0)}, |
| 534 | ) | 534 | ) |
| 535 | + | ||
| 536 | + | ||
| 537 | +@torch.inference_mode() | ||
| 538 | +def test_per_token_cast_compile_cache_layout(): | ||
| 539 | + # Compile quantization only; cache APIs themselves are outside this change. | ||
| 540 | + from tests.testing.numeric import assert_equal | ||
| 541 | + | ||
| 542 | + x = torch.randn((31, 512), dtype=torch.bfloat16, device=get_device()) | ||
| 543 | + slots = torch.arange(31, dtype=torch.int32, device=x.device) | ||
| 544 | + cache = torch.empty((4, 32, 1, 528), dtype=torch.uint8, device=x.device) | ||
| 545 | + | ||
| 546 | + def quantize(value): | ||
| 547 | + return per_token_cast(value, 'e4m3', 32, round_sf=True, use_packed_ue8m0=True) | ||
| 548 | + | ||
| 549 | + compiled = torch.compile(quantize, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 550 | + for step in range(3): | ||
| 551 | + if step == 1: | ||
| 552 | + x.neg_().add_(0.125) | ||
| 553 | + elif step == 2: | ||
| 554 | + x = torch.randn_like(x) | ||
| 555 | + cache.fill_(0x5A) | ||
| 556 | + expected_data, expected_sf = torch_cast( | ||
| 557 | + x, | ||
| 558 | + 'e4m3', | ||
| 559 | + (1, 32), | ||
| 560 | + round_sf=True, | ||
| 561 | + use_packed_ue8m0=True, | ||
| 562 | + ) | ||
| 563 | + data, scale = compiled(x) | ||
| 564 | + assert_equal(data, expected_data) | ||
| 565 | + assert_equal(scale, expected_sf) | ||
| 566 | + store_kv_cache(x.unsqueeze(1), cache, slots, use_fp8_rope=True, num_per_channels=32) | ||
| 567 | + raw = cache.view(4, 32 * 528) | ||
| 568 | + data_bytes_per_block = 32 * 512 | ||
| 569 | + stored_data = raw[:, :data_bytes_per_block].view(4, 32, 512) | ||
| 570 | + stored_sf = raw[:, data_bytes_per_block:].view(4, 32, 16) | ||
| 571 | + assert_equal(stored_data[0, :31], data.view(torch.uint8)) | ||
| 572 | + assert_equal(stored_sf[0, :31], scale.view(torch.uint8)) | ||
| 573 | + assert torch.all(stored_data[0, 31] == 0x5A) | ||
| 574 | + assert torch.all(stored_sf[0, 31] == 0x5A) | ||
| 575 | + del data, scale | ||
| @@ -409,3 +409,30 @@ def test_norm_fp8_cast_benchmark( | |||
| 409 | time_us=t_us, | 409 | time_us=t_us, |
| 410 | extras={'bandwidth_gbs': io_bytes / 1e9 / (t_us / 1e6)}, | 410 | extras={'bandwidth_gbs': io_bytes / 1e9 / (t_us / 1e6)}, |
| 411 | ) | 411 | ) |
| 412 | + | ||
| 413 | + | ||
| 414 | +@torch.inference_mode() | ||
| 415 | +def test_norm_per_token_cast_compile_composition(): | ||
| 416 | + from tests.quant.cast_ref import cast | ||
| 417 | + from tests.testing.numeric import assert_equal | ||
| 418 | + | ||
| 419 | + x = torch.randn((31, 512), dtype=torch.bfloat16, device=get_device()) | ||
| 420 | + weight = torch.randn((512,), dtype=torch.bfloat16, device=x.device) | ||
| 421 | + | ||
| 422 | + def model(value, norm_weight): | ||
| 423 | + normalized = rms_norm_ref(value, norm_weight, 1e-6) | ||
| 424 | + return per_token_cast(normalized, 'e4m3', 32, round_sf=True) | ||
| 425 | + | ||
| 426 | + compiled = torch.compile(model, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 427 | + for step in range(3): | ||
| 428 | + if step == 1: | ||
| 429 | + x.neg_().add_(0.125) | ||
| 430 | + weight.neg_() | ||
| 431 | + elif step == 2: | ||
| 432 | + x = torch.randn_like(x) | ||
| 433 | + weight = torch.randn_like(weight) | ||
| 434 | + expected = cast(rms_norm_ref(x, weight, 1e-6), 'e4m3', (1, 32), round_sf=True) | ||
| 435 | + outputs = compiled(x, weight) | ||
| 436 | + for actual, reference in zip(outputs, expected, strict=True): | ||
| 437 | + assert_equal(actual, reference) | ||
| 438 | + del outputs, actual | ||
| @@ -164,3 +164,28 @@ def test_cast_back_benchmark(benchmark_timer, benchmark_record, params): | |||
| 164 | time_us=t_us, | 164 | time_us=t_us, |
| 165 | bandwidth_gbs=num_bytes / t_us / 1e3, | 165 | bandwidth_gbs=num_bytes / t_us / 1e3, |
| 166 | ) | 166 | ) |
| 167 | + | ||
| 168 | + | ||
| 169 | +@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1']) | ||
| 170 | +@torch.inference_mode() | ||
| 171 | +def test_per_token_cast_compile_then_cast_back(fmt): | ||
| 172 | + x = torch.randn((31, 512), dtype=torch.bfloat16, device=get_device()) | ||
| 173 | + | ||
| 174 | + def model(value): | ||
| 175 | + sf = cann_ops_tilelang.per_token_cast_with_sf_only(value, fmt, 32, round_sf=True) | ||
| 176 | + data = cann_ops_tilelang.per_token_cast_with_precomputed_sf(value, fmt, 32, sf, round_sf=True) | ||
| 177 | + return data, sf | ||
| 178 | + | ||
| 179 | + compiled = torch.compile(model, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 180 | + for step in range(3): | ||
| 181 | + if step == 1: | ||
| 182 | + x.neg_().add_(0.125) | ||
| 183 | + elif step == 2: | ||
| 184 | + x = torch.randn_like(x) | ||
| 185 | + reference = cast(x, fmt, (1, 32), round_sf=True) | ||
| 186 | + data, sf = compiled(x) | ||
| 187 | + assert_equal(data, reference[0]) | ||
| 188 | + assert_equal(sf, reference[1]) | ||
| 189 | + restored = cann_ops_tilelang.per_token_cast_back((data, sf), 'fp32', num_per_channels=32) | ||
| 190 | + assert_equal(restored, cast_back(reference, 'fp32', (1, 32))) | ||
| 191 | + del data, sf | ||
| @@ -1,3 +1,4 @@ | |||
| 1 | +import logging | ||
| 1 | import math | 2 | import math |
| 2 | import os | 3 | import os |
| 3 | 4 | ||