Skip to content

[XPU][FP8] Reuse block MoE prefill descriptor buffers - #573

Merged
gc-fu merged 2 commits into
mainfrom
pr/fp8-block-moe-prefill-buffer-cache
Jul 28, 2026
Merged

[XPU][FP8] Reuse block MoE prefill descriptor buffers#573
gc-fu merged 2 commits into
mainfrom
pr/fp8-block-moe-prefill-buffer-cache

Conversation

@gc-fu

@gc-fu gc-fu commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Summary

  • reuse tile_experts and tile_rows tensors across block-scaled MoE prefill calls
  • fill unused expert descriptors inside the descriptor-builder kernel, removing the separate tensor-fill submission
  • grow cached tensors only when the required tile capacity increases
  • isolate cached buffers by host thread, XPU device, and SYCL queue so concurrent streams cannot overwrite one another

Stream safety

A single thread-local buffer is not sufficient because one host thread can enqueue work on multiple XPU streams while prior work remains asynchronous. This PR keeps a persistent cache entry for each (device_index, sycl::queue) pair. Calls on the same stream reuse ordered storage; calls on different streams use distinct tensors.

Scope

This is a kernel-only performance optimization following #566 and the schema-only #572. The existing operator and framework dispatch are unchanged.

Testing

Not run as part of PR creation.

@gc-fu

gc-fu commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Validation completed on XPU 0,1.

Code state:

  • Framework: 4d20a859d
  • Kernel PR head: e64ae59
  • oneCCL used by libtorch.so: /opt/intel/oneapi/ccl/2021.15/lib/libccl.so.1

Accuracy

GSM8K configuration: 30 rounds x 100 questions, 5-shot, temperature 0, min_tokens=2.

Workload Result Mean accuracy Range Invalid NaN Device errors
Offline per-block FP8 30/30 PASS 98.37% 97%-99% 0 0 0
Online per-tensor FP8 30/30 PASS 96.93% 96%-98% 0 0 0

Per-block performance

Same accuracy-warmed service, exact 32k input / 2k output, one round:

Batch Output tok/s TTFT (ms) TPOT (ms) E2E (ms)
1 53.9714 4212.869 16.4791 37945.654
4 117.6631 10836.386 28.6511 69485.104

The random numerical correctness test for the FP8 block-scale MoE prefill kernel passed. The exhaustive E4M3 encoding test still exposes the existing raw 0x80 negative-zero issue in the dense path (expected -0.0, actual -0.5). This is outside the buffer-cache changes in this PR and belongs to the deferred full FP8/GDN hardening work.

The wheel build completed successfully after one full retry. The first attempt hit a transient llvm-foreach segmentation fault while device-linking the unchanged esimd_topk_v2 extension; it did not reproduce on retry.

Conclusion: no per-block accuracy or performance regression was observed, and no cross-workload accuracy regression was observed for online per-tensor FP8. This PR is ready to merge based on the completed validation.

@gc-fu
gc-fu merged commit 51d4a95 into main Jul 28, 2026
@gc-fu
gc-fu deleted the pr/fp8-block-moe-prefill-buffer-cache branch July 28, 2026 01:36
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.

1 participant