118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
"""comfy-kitchen registration for the gfx1103 native Triton W4A4 path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import sys
|
|
|
|
import torch
|
|
|
|
from comfy_kitchen.constraints import (
|
|
DivisibleBy,
|
|
ExactDims,
|
|
FunctionConstraints,
|
|
MinDims,
|
|
ParamConstraint,
|
|
)
|
|
from comfy_kitchen.registry import registry
|
|
|
|
from .native_int4 import is_gfx1103_target, triton_convrot_w4a4_linear
|
|
|
|
|
|
BACKEND_NAME = "gfx1103_int4_triton"
|
|
|
|
|
|
class _Int4LinearDtypeConstraint(ParamConstraint):
|
|
"""Accept the W4A4 backend only for an explicitly native INT4 call.
|
|
|
|
comfy-kitchen also uses ``linear_dtype='int8'`` for a deliberate mixed
|
|
precision tier. Returning a failed constraint for that value lets the
|
|
registry continue to an INT8-capable backend instead of silently treating
|
|
it as native W4A4.
|
|
"""
|
|
|
|
def check_dtype(self, value) -> bool:
|
|
return value == "int4"
|
|
|
|
|
|
class _ConvRotGroupConstraint(ParamConstraint):
|
|
def check_dtype(self, value) -> bool:
|
|
return type(value) is int and value in (16, 64, 256)
|
|
|
|
|
|
class _QuantGroupConstraint(ParamConstraint):
|
|
def check_dtype(self, value) -> bool:
|
|
return type(value) is int and value == 64
|
|
|
|
|
|
def convrot_w4a4_linear(
|
|
x: torch.Tensor,
|
|
qweight: torch.Tensor,
|
|
wscales: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
convrot_groupsize: int = 256,
|
|
quant_group_size: int = 64,
|
|
linear_dtype: str = "int4",
|
|
) -> torch.Tensor:
|
|
return triton_convrot_w4a4_linear(
|
|
x,
|
|
qweight,
|
|
wscales,
|
|
bias,
|
|
convrot_groupsize=convrot_groupsize,
|
|
quant_group_size=quant_group_size,
|
|
linear_dtype=linear_dtype,
|
|
output_dtype=x.dtype,
|
|
)
|
|
|
|
|
|
def _constraints() -> dict[str, FunctionConstraints]:
|
|
floats = frozenset({torch.float16, torch.bfloat16})
|
|
return {
|
|
"convrot_w4a4_linear": FunctionConstraints(
|
|
params={
|
|
"x": ParamConstraint(dtypes=floats, shape_rules=(MinDims(2),)),
|
|
"qweight": ParamConstraint(
|
|
dtypes=frozenset({torch.int8}),
|
|
shape_rules=(ExactDims(2), DivisibleBy(dim=1, factor=32)),
|
|
),
|
|
"wscales": ParamConstraint(
|
|
dtypes=frozenset({torch.float16, torch.bfloat16, torch.float32}),
|
|
shape_rules=(ExactDims(1),),
|
|
),
|
|
"bias": ParamConstraint(
|
|
dtypes=frozenset({torch.float16, torch.bfloat16, torch.float32})
|
|
),
|
|
"convrot_groupsize": _ConvRotGroupConstraint(dtypes=frozenset({int})),
|
|
"quant_group_size": _QuantGroupConstraint(dtypes=frozenset({int})),
|
|
"linear_dtype": _Int4LinearDtypeConstraint(dtypes=frozenset({str})),
|
|
},
|
|
# ZLUDA exposes the AMD device through torch's CUDA device type.
|
|
default_devices=frozenset({"cuda"}),
|
|
)
|
|
}
|
|
|
|
|
|
def register_kitchen_backend() -> bool:
|
|
"""Register ahead of stock backends without replacing any of them."""
|
|
if not is_gfx1103_target():
|
|
logging.info("INT4 Fast: gfx1103 Triton target not active; kitchen backend not registered.")
|
|
return False
|
|
|
|
registry.register(
|
|
name=BACKEND_NAME,
|
|
module=sys.modules[__name__],
|
|
capabilities=_constraints(),
|
|
)
|
|
# comfy-kitchen 0.2.x has a setter but no public priority getter. Preserve
|
|
# every existing entry and prepend only our narrowly-scoped backend.
|
|
current_priority = list(registry._priority)
|
|
registry.set_priority(
|
|
[BACKEND_NAME] + [name for name in current_priority if name != BACKEND_NAME]
|
|
)
|
|
logging.info(
|
|
"INT4 Fast: registered comfy-kitchen backend %s for native gfx1103 W4A4.",
|
|
BACKEND_NAME,
|
|
)
|
|
return True
|