jLdZddlZddlZddlmZddlmZddlmZeZ ej dej dej diZ d Zejd Zdd d d ZejdZejdZdZdd d d d ddZej fd dd ddZejddS)a{Common eligibility and family ordering for activation kernels. Dispatch never caches tensors, modules, bound methods or shape-specific choices. An operation supplies only the implementations that satisfy its arithmetic and layout contract; the existing user policy orders those implementations. Device eligibility lives here, once: a WanGP-owned kernel runs on every NVIDIA GPU whose instructions its arithmetic needs (DTYPE_FLOORS), not only where it was first timed. Triton compiles each kernel for the running GPU; the kernels carry no architecture-specific code, and their precision against PyTorch was measured per layer and end to end on compute capabilities 8.9 and 12.0 (specs: H3_FL2VA_ACCELERATION_AUDIT_20261021*.md, H3_VAE_SPEED_20261009.md). An operation whose kernel needs more than its dtype floor says so (`capability=`); one validated on HIP admits it (`hip=True`); PACKED_FMA_CAPABILITY marks the one architecture whose compiler changes a bit-exact kernel's rounding. Performance-only restrictions (a kernel that runs everywhere but only beats cuDNN on RTX 50) stay with the kernel that measured them, built on `capability()`. N)tqdm) cuda_graphs) kernel_policyr)) rctjjdStj|tjn|S)z{(major, minor) of an NVIDIA CUDA device, None on HIP. Device capabilities are immutable; cache metadata, never activations.N)torchversionhipcudaget_device_capabilitycurrent_device) device_indexs ?C:\Users\johnn\OneDrive\Documents\AI\shared\kernels\dispatch.py capabilityr%sC }$t : + +rc|tjkpJtjjdup7tj|t tjkS)zBF16 arithmetic needs native conversion instructions on NVIDIA GPUs (HIP: Triton's own conversions). Self-contained, so that clearing its cache alone re-reads the device (the LM engine and FlashVSR tests patch the capability around it).N)r bfloat16rrrrr)rrs rtriton_dtype_supportedr%:s^ EN " Pem&7t&C PuzGgGghnGoGosAFAOtPHPPrcVt|}|duo|dtdkS)zWhether the compiler of this device fuses a packed multiply and add despite enable_fp_fusion=False (PACKED_FMA_CAPABILITY): a kernel bit-identical to PyTorch elsewhere is one rounding apart here, so it counts as approximate on this device only.Nr)rPACKED_FMA_CAPABILITY)rr s r packed_fmar(As0 | $ $E   Eq-B1-E!EErct|trt|dkrtd|S)zConsume an owned activation. Callers must drop their other references. Residuals, parameters, rotary tables and caches are borrowed explicitly, not put in this holder. In-place APIs ending in '_' borrow their destination. rz(Expected a one-element activation holder) isinstancelistlen ValueErrorpop)x_lists rtaker0IsD fd # #Es6{{a'7'7CDDD ::<<r) wgpkernelskitchentriton owned_tritonc Ntj|p|||p|d}|tjsh|dkr|sdn|pd}||tjf}|t vrBt |tj d|d|dtjd|S) N)r1r2r3r1z WanGP Tritonpytorchz [Kernels] z: z (z).) rfirstr compiler is_compiling precision _announcedaddrwrite) operationr1r2r3r4familyimplementationkeys rchooserBTs 0Jl-4@V,"X"XYYFU^%@%@%B%B+1\+A+A*+AZ`Zmdm)@@ j NN3    J^I^^^^=CZ^^^ _ _ _ Mr)dtypes approximaterrc"tjsStjs@t |tjus%|jr|j|vs|rtj sdSt|j |j||S)a7Whether a WanGP kernel may replace the PyTorch path for x: no autograd (inference mode, or grad disabled as wgp.py runs every generation: the kernels never touch autograd, and pipelines under no_grad alone, such as Flux, Qwen Image and Wan, were running the PyTorch paths while their blocks measured the kernels under inference mode), no compilation, a CUDA tensor of one of dtypes on a device the kernel supports (device_supported), and the user's precision policy when the kernel is approximate (a bool, possibly decided per device: see packed_fma).Fr) r r8r9is_grad_enabledrTensoris_cudarrallow_approximater"r)xrCrDrrs rcuda_inferencerKbs ##%%)>)@)@Awwel**!)*qwf?T?T@U$1$C$E$E@Uu AHag@PVY Z Z ZZrc2tjtjfSN)rr:familiesrrrPqs-"9=;Q!RrrM)__doc__ functoolsr rmmgprrsetr;r$float16float32rr'cacherr"r%r(r0rBrK register_keyrOrrrZs$  SUU  vu}fU   sss9=%      PPP  FFF eZ_     "' 1uW[af [ [ [ [ [ RRSSSSSr