wehub-resource-sync
83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License").
|
|
# You may not use this file except in compliance with the License.
|
|
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
import torch
|
|
|
|
try:
|
|
from . import longlive_kv_dequant_cuda # noqa: F401
|
|
except ImportError:
|
|
import longlive_kv_dequant_cuda # noqa: F401
|
|
|
|
|
|
def _dtype_to_code(dtype: torch.dtype) -> int:
|
|
if dtype == torch.bfloat16:
|
|
return 0
|
|
if dtype == torch.float16:
|
|
return 1
|
|
if dtype == torch.float32:
|
|
return 2
|
|
raise ValueError(f"Unsupported fused KV dequant dtype: {dtype}")
|
|
|
|
|
|
def scale_rule_to_fp4_limits(scale_rule) -> tuple[float, float]:
|
|
"""Return the dequant denominator limits used by FourOverSix ScaleRule."""
|
|
if hasattr(scale_rule, "max_allowed_e2m1_value") and hasattr(
|
|
scale_rule, "max_allowed_e4m3_value",
|
|
):
|
|
return (
|
|
float(scale_rule.max_allowed_e2m1_value()),
|
|
float(scale_rule.max_allowed_e4m3_value()),
|
|
)
|
|
|
|
normalized = str(scale_rule).lower()
|
|
if "." in normalized:
|
|
normalized = normalized.rsplit(".", 1)[-1]
|
|
normalized = normalized.strip().strip("\"'")
|
|
|
|
if normalized == "static_4":
|
|
return 4.0, 448.0
|
|
if normalized == "static_6":
|
|
return 6.0, 448.0
|
|
if normalized in {"mse", "mae", "l1_norm", "abs_max"}:
|
|
return 6.0, 256.0
|
|
|
|
raise ValueError(f"Unsupported FP4 scale_rule: {scale_rule}")
|
|
|
|
|
|
def dequantize_kv_cache_fp4(
|
|
values: list[torch.Tensor],
|
|
scale_factors: list[torch.Tensor],
|
|
amax: list[torch.Tensor],
|
|
*,
|
|
num_heads: int,
|
|
block_token_size: int,
|
|
dtype: torch.dtype,
|
|
e2m1_max: float | None = None,
|
|
e4m3_max: float | None = None,
|
|
scale_rule=None,
|
|
) -> torch.Tensor:
|
|
"""Dequantize multiple AR KV-cache chunks with one CUDA launch."""
|
|
if e2m1_max is None or e4m3_max is None:
|
|
if scale_rule is None:
|
|
raise ValueError(
|
|
"Either e2m1_max/e4m3_max or scale_rule must be provided.",
|
|
)
|
|
e2m1_max, e4m3_max = scale_rule_to_fp4_limits(scale_rule)
|
|
|
|
return torch.ops.longlive_kernels.dequantize_kv_cache_fp4.default(
|
|
values,
|
|
scale_factors,
|
|
amax,
|
|
num_heads,
|
|
block_token_size,
|
|
_dtype_to_code(dtype),
|
|
e2m1_max,
|
|
e4m3_max,
|
|
)
|