submission 513577

"""
V2 variant: Same CUDA graph approach but with Triton quant instead of HIP quant.
Try this if HIP quant is slower or unavailable.
"""
import torch
from task import input_t, output_t
import aiter
from aiter import QuantType, dtypes

_quant = aiter.get_triton_quant(QuantType.per_1x32)
_graphs: dict = {}
_no_graph: set = set()


def custom_kernel(data: input_t) -> output_t:
    A, _, _, B_shuffle, B_scale_sh = data
    key = (A.data_ptr(), B_shuffle.data_ptr(), B_scale_sh.data_ptr())

    cached = _graphs.get(key)
    if cached is not None:
        cached[0].replay()
        return cached[1]

    A_q, A_s = _quant(A, shuffle=True)
    out = aiter.gemm_a4w4(
        A_q, B_shuffle, A_s, B_scale_sh,
        dtype=dtypes.bf16, bpreshuffle=True,
    )

    if key not in _no_graph:
        try:
            torch.cuda.synchronize()
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g):
                _aq, _as = _quant(A, shuffle=True)
                _out = aiter.gemm_a4w4(
                    _aq, B_shuffle, _as, B_scale_sh,
                    dtype=dtypes.bf16, bpreshuffle=True,
                )
            _graphs[key] = (g, _out)
        except Exception:
            _no_graph.add(key)

    return out