[perf] FastH3 two-Spark: gather packed FP4 inputs before head projections

Opt-in h3_vsa_fp4.sp_input_gather (default false; requires sp_projection_overlap) gathers the
already-packed FP4 inputs and rotary tables, then computes each rank's assigned Q/K/V heads locally;
the gate projection overlaps the smaller transfer. Supported on the measured GB10/cuDNN shape only;
other shapes keep the existing exchange. Enabled in the Spark-pair example.
T2AV on 2 Sparks: 19.35 -> 18.76 s/clip (merge gate check, same SP reduction order rules).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Junda Chen
2026-10-24 02:16:48 +00:00
co-authored by Claude Opus 5.5
parent 6397822a5b
commit be8a13b749
7 changed files with 308 additions and 9 deletions
@@ -162,6 +162,14 @@ projection is packed, overlapping traffic with the remaining projection GEMMs. I
QKV/gate projections with shared input quantization. Other configurations use the ordinary
head-group exchange. Packing arithmetic and FP8 transport precision are unchanged.
The optional `h3_vsa_fp4.sp_input_gather: true` gathers the existing packed FP4 inputs and
rotary tables before computing each rank's assigned Q/K/V heads. The gate projection overlaps
that smaller input transfer. It defaults to `false`, requires `sp_projection_overlap: true`,
and currently supports the measured GB10/cuDNN shape: 7430 local rows, 5376 input channels,
56 heads and bias-free NVFP4 QKV projections. Other shapes use the existing exchange path.
All projections, normalization, rotary operations and FP8 packet rounding remain present;
no steps or tokens are omitted. Validate generated output when changing this GEMM shape.
The optional `h3_vsa_fp4.sp_pooled_key_mean: true` reuses tile statistics for key smoothing
on the grouped tile-64 path with FP8 prefix attention and `smooth_k: true`. It defaults to
`false`. All key rows still contribute, including zero padding in the normalization length,