mstar.model.z_image.components.qk_norm_rope#
Z-Image’s per-head q/k RMSNorm and rotary in one pass, for the compiled transformer.
Inductor splits ScaledRMSNorm + the rotary product into a variance reduction per q and k, an
apply pass and a layout copy (~42 ms of a ~800 ms 1024^2 request on an H200). This kernel reads
each 128-wide head row of q and k once from the fused qkv projection, and writes the rotated
rows contiguous [B, S, H, D], rounding to bf16 where the eager modules do: after the
normalization, after * weight, and after the fp32 rotary product.
Functions
|