[ET-VK][quantized] Store dq8ca per-token zero-point as fp32#20491
[ET-VK][quantized] Store dq8ca per-token zero-point as fp32#20491SS-JIA wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20491
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ef9eec1 with merge base 1227757 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
This PR needs a
|
…type Pull Request resolved: #20491 The per-token dynamic-activation-quant (`dq8ca`) zero-point image must be bound in the shader with the same dtype the tensor was allocated with; a binding-vs-allocation dtype mismatch corrupts the per-token zero-point. The allocation dtype differs by export path: standard `export_llama -qmode 8da4w` models (e.g. Qwen3-0.6B) serialize the zero-point as `int8`, while the Llama4-mini TISO backbone (torchao `per_token_dynamic_quant` / `Int8DynamicActivationIntxWeightConfig` with an explicit fp32 `zero_point_dtype`) serializes it as float, which `vulkan_graph_builder.get_effective_dtype` downcasts to `half` under `force_fp16`. A single fixed binding dtype cannot satisfy both paths. Binding the zero-point as `int8` (`rgba8i`) corrupts the float-allocated TISO zero-point on ARM Mali (Valhall) -- negative values come back as garbage, garbling the 8da4w TTS backbone. Conversely, binding it as the codegen `DTYPE` (matching the scale's float dtype) corrupts the int8-allocated zero-point: under fp16 inference the `rgba8i` image is read and written as `rgba16f`, saturating the per-token zero-point to the int8 floor/ceiling and garbling standard fp16 8da4w models such as Qwen3-0.6B. This change makes the zero-point binding a codegen variant so it always matches the tensor's allocation. A new `ZP_DTYPE_MODE` axis emits two variants of every dq8ca shader that binds the per-token zero-point: `zpint8` (binding declared `int8`, an `rgba8i` integer image) and `zpinherit` (binding declared with the codegen `DTYPE`, inheriting the inference float dtype to match the scale -- `rgba32f`, or `rgba16f` under `USE_VULKAN_FP16_INFERENCE`). The C++ shader pickers select the variant from `graph.dtype_of(zero_point)` (`kChar` -> `zpint8`; `kHalf` / `kFloat` -> `zpinherit`), so the shader binding matches the tensor's allocation regardless of how the model was exported. The shared read helper is unchanged: `ivec4(texelFetch(t_int8_input_zps, ...))` already reads both an integer image (identity) and a float image (exact truncation of the integer-valued zero-point in `[-128, 127]`). Affected shaders: `choose_qparams_per_row` (writes the zero-point, storing `ivec4` or `VEC4_T` per variant), `quantize_and_pack_4h4w_with_group_sums`, `linear_dq8ca_q4gsw_tiled`, and the dq8ca `linear_q4gsw_coop` variants (read the zero-point). This fixes the fp16 8da4w regression for standard int8 zero-point exports while preserving the float zero-point path that the TISO backbone and the original Mali fix depend on. Only the runtime shader binding changes, so existing `.pte` files are handled correctly with no re-export. Authored with Claude Code. ghstack-source-id: 397279874 @exported-using-ghexport Differential Revision: [D109595977](https://our.internmc.facebook.com/intern/diff/D109595977/)
|
@claude review this code |
|
Claude finished @metascroy's task in 8m 41s —— View job Code Review
I traced the full dq8ca per-token path (choose_qparams → quantize_and_pack → linear) plus the shared loader and the coop binding. The fix is correct and well-targeted. Summary and a few minor notes below. Correctness — looks right ✅
Minor notes (non-blocking)
Overall: focused, correct fix with a sensible codegen abstraction. The notes above are polish, not blockers. |
2 participants
Stack from ghstack (oldest at bottom):
The per-token dynamic-activation-quant (
dq8ca) zero-point was corrupted by a tensor-allocation vs shader-access dtype mismatch. The per-token zero-point tensor is created with a float dtype -- fp32, or fp16 underUSE_VULKAN_FP16_INFERENCE-- so its backing image uses a float texel format (rgba32f/rgba16f). But the shader declared and accessed that image with an integer dtype (int8, an integer image formatrgba8i). Reading a float-format image through an integer-format binding is the bug. On ARM Mali (Valhall) GPUs this mismatch corrupted the per-token zero-points: negative zero-points came back as garbage (-kread as-2^23 - k), driving the quantized activation to the int8 floor, the per-group sums to-4096, and the GEMM output to garbage, producing garbled, runaway generation for 8da4w models (e.g. the Llama4-mini TISO TTS backbone on Mali-G715/G710). Adreno happened to tolerate the same mismatch and read correct values, so the corruption was Mali-specific even though the mismatch itself is general.The per-token zero-point is serialized as fp32 by torchao design:
Int8DynamicActivationIntxWeightConfig(8da4w) uses asymmetric per-token activation quant with an explicit fp32zero_point_dtype. Decoding the serialized.pteconfirms the zero-point tensor is FLOAT32, and (like the scale) it is stored in a texture as anrgba32ftexel -- neverrgba8i. The float allocation is the truth; the int8 shader access was the mismatched side.The fix is to declare, store, and read the per-token zero-point as fp32 across the dq8ca qparams shaders, so the shader access dtype matches the tensor's allocation dtype and the texture is read as the
rgba32fimage it actually is. The zero-point value is integer-valued (nudged to[-128, 127]), so fp32 represents it exactly and the consumer'sint(zp)conversion for the integer dequant-correction is lossless. This touches the dq8ca qparams shaders --choose_qparams_per_row,quantize_and_pack_4h4w_with_group_sums,linear_dq8ca_q4gsw_tiled, the sharedlinear_int8_input_scales_zps_loadhelper, and thelinear_q4gsw_coopvariant (whose zero-point binding only matches the descriptor-set layout and is never read) -- plus a documentation comment inChooseQParams.cpp.Because the per-token qparams remain in texture storage (unchanged from before) and only the zero-point dtype changes, this is a pure runtime shader fix: existing texture-qparams 8da4w
.ptefiles are corrected without re-export, since the texture already bakes the zero-point asrgba32fand the shader now reads it as such.Authored with Claude Code.
Differential Revision: D109595977