Skip to content

Add accumulation precision, TPU v6e tests, and SQNR benchmarks for multi-pass FP8 in QWIX - #378

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_982908980
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_982908980

Conversation

@copybara-service

Copy link
Copy Markdown

Add accumulation precision, TPU v6e tests, and SQNR benchmarks for multi-pass FP8 in QWIX

  • Add unit test evaluating multi-pass FP8 across accumulation regimes:
    • FP32 accumulation (unconstrained mantissa)
    • Reference software emulation with output rounded to BF16
    • Internal BF16 accumulator (preferred_element_type=bfloat16), demonstrating exponent alignment truncation of the 2^-4 residual passes
  • Add TPU v6e test (multipass_dot_tpu_test) with requires-ghostfish tag to verify execution on Ghostfish hardware.
  • Add comprehensive SQNR benchmark summary table across multi-pass FP8 strategies contrasting CPU software accumulation against physical TPU hardware accumulation (512x512, channel-wise / tensor scaling).

Multi-Pass FP8 SQNR Benchmark Table (512x512, Channel-wise / Tensor Scaling)

Strategy GEMMs Compute DType Execution Target Accumulation Mode SQNR (dB) Rel Error (%)
Baseline BF16 (Reference) 1 bfloat16 TPU / CPU Output Rounded to BF16 55.62 dB 0.17%
Pure FP8 (1-Pass) 1 float8_e4m3fn CPU Reference FP32 Accumulation (Rounded to BF16) 28.43 dB 3.79%
Pure FP8 (1-Pass) 1 float8_e4m3fn Physical TPU Hardware Accumulator 28.39 dB 3.80%
Triangular Multi-Pass FP8 3 float8_e4m3fn CPU Reference FP32 Accumulation (Rounded to BF16) 54.84 dB 0.18%
Triangular Multi-Pass FP8 3 float8_e4m3fn Physical TPU Hardware Accumulator 47.48 dB 0.42%
Full Cross Multi-Pass FP8 4 float8_e4m3fn CPU Reference FP32 Accumulation (Rounded to BF16) 55.59 dB 0.17%
Full Cross Multi-Pass FP8 4 float8_e4m3fn Physical TPU Hardware Accumulator 47.60 dB 0.42%

…lti-pass FP8 in QWIX

- Add unit test evaluating multi-pass FP8 across accumulation regimes:
  - FP32 accumulation (unconstrained mantissa)
  - Reference software emulation with output rounded to BF16
  - Internal BF16 accumulator (preferred_element_type=bfloat16), demonstrating exponent alignment truncation of the 2^-4 residual passes
- Add TPU v6e test (multipass_dot_tpu_test) with requires-ghostfish tag to verify execution on Ghostfish hardware.
- Add comprehensive SQNR benchmark summary table across multi-pass FP8 strategies contrasting CPU software accumulation against physical TPU hardware accumulation (512x512, channel-wise / tensor scaling).

### Multi-Pass FP8 SQNR Benchmark Table (512x512, Channel-wise / Tensor Scaling)

| Strategy | GEMMs | Compute DType | Execution Target | Accumulation Mode | SQNR (dB) | Rel Error (%) |
| :---| :---: | :---| :---| :---| :---: | :---: |
| Baseline BF16 (Reference) | 1 | bfloat16 | TPU / CPU | Output Rounded to BF16 | 55.62 dB | 0.17% |
| Pure FP8 (1-Pass) | 1 | float8_e4m3fn | CPU Reference | FP32 Accumulation (Rounded to BF16) | 28.43 dB | 3.79% |
| Pure FP8 (1-Pass) | 1 | float8_e4m3fn | Physical TPU | Hardware Accumulator | 28.39 dB | 3.80% |
| Triangular Multi-Pass FP8 | 3 | float8_e4m3fn | CPU Reference | FP32 Accumulation (Rounded to BF16) | 54.84 dB | 0.18% |
| Triangular Multi-Pass FP8 | 3 | float8_e4m3fn | Physical TPU | Hardware Accumulator | 47.48 dB | 0.42% |
| Full Cross Multi-Pass FP8 | 4 | float8_e4m3fn | CPU Reference | FP32 Accumulation (Rounded to BF16) | 55.59 dB | 0.17% |
| Full Cross Multi-Pass FP8 | 4 | float8_e4m3fn | Physical TPU | Hardware Accumulator | 47.60 dB | 0.42% |

PiperOrigin-RevId: 982908980
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants