Module innerq_driver

Module innerq_driver 

Source
Expand description

Host-side driver for TurboQuant+ InnerQ per-channel K equalization.

Triggers via TURBO_INNERQ=N env var (N = calibration token count). The kernel-side state lives in kernels/gb10/common/tq_plus_innerq_apply.cu as __device__ globals inside namespace tq_plus — deliberately in the SAME translation unit (= same PTX module; Atlas has no -rdc device linking) as the apply/accumulate kernels that read it. This driver manipulates that state directly via the CUDA Driver API:

cuModuleGetGlobal_v2 → device pointer for each symbol cuMemcpyHtoDAsync_v2 / cuMemcpyDtoHAsync_v2 → push/pull state

Two-phase operation:

  1. start() — zero counters, set d_innerq_calibrating = 1.
  2. maybe_finalize() — read d_innerq_count; once it crosses target_tokens, read d_innerq_sq_accum, compute per-channel scale + scale_inv, upload, set d_innerq_active = 1.

Math identity: <Q/s, s·K> = <Q, K> — the kernel-side apply pass multiplies Q by scale_inv pre-WHT and K by scale post-WHT, leaving attention dot products unchanged while smoothing K variance across channels.

Structs§

InnerQDriver