Repository navigation
Conversation
- Accept caller-owned A8W8 outputs and optionally reuse A4W4 scratch. - Reuse the largest tuned bpreshuffle row only above the table and chunk FlyDSL launches at their row limit. - Keep KDA prefill tuning and dispatch tests with the GEMM implementation.
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
|
Closing: the large-M reuse heuristic (reuse the largest tuned bpreshuffle row up to 3x its M) slowed down 9 of 43 gfx950 tuned groups in a sweep, the worst at 0.17x (DSv3 N=2112 K=7168). Kimi-K3's missing M values are being added as tuned rows instead (see #6410). The caller-owned |
Motivation
A8W8 bpreshuffle GEMM allocates an output even when the caller already owns a suitable destination, and a row count above the tuned table misses the available configuration.
Technical Details
A8W8 bpreshuffle GEMM allocates an output even when the caller already owns a suitable destination, and a row count above the tuned table misses the available configuration. Add a validated optional out buffer, opt-in A4W4 scratch reuse, and reuse of the largest matching tuned M only when the requested M exceeds it. Chunk FlyDSL launches at their row limit and slice per-row activation scales with each chunk. Keep the Kimi-K3 FP8 tuning rows and large-M dispatch tests with this change.
Branch
pr/aiter-03-gemm-output-and-large-m, head10ae462f18b9e352b210de539146f4657c26062a; intended feature basee4ceb9dd95e12595dc9d87a06eaef79dfd6cdb8c. 6 files changed, 395 insertions(+), 23 deletions(-).Changed files in this slice
aiter/configs/model_configs/a8w8_bpreshuffle_tuned_gemm_kimik3.csvaiter/configs/model_configs/a8w8_bpreshuffle_untuned_gemm_kimik3.csvaiter/configs/model_configs/dsv3_a8w8_bpreshuffle_tuned_gemm.csvaiter/ops/gemm_op_a4w4.pyaiter/ops/gemm_op_a8w8.pyop_tests/test_gemm_a8w8_bpreshuffle_large_m.pyTest Plan
Validation performed and its limits are listed below. Remaining before marking ready: Exact-head large-M numerical equivalence across chunk boundaries and graph replay with caller-owned buffers.
Test Result
op_tests/test_gemm_a8w8_bpreshuffle_large_m.pypassed on combined AITER: dispatch above the largest tuned M and fallback below the table are exercised. This was part of a 4-test run, completed in 4.45 s. Standalone direct probes below use adjusted SGLang targeta1e96661f4937e7ee4fa811c2f28d26297db82e3with AITERd78f1d4751fa.Direct operator probe:
{"same_stream_reuses":true,"different_stream_isolated":true}.Direct operator probe:
{"caller_output_reused":true,"exact_equality":true,"invalid_shape_dtype_device":["expected ValueError","expected ValueError","expected ValueError"]}.No isolated GEMM throughput or Kimi-K3 E2E delta is measured for this head. A
[16384,7168]BF16 output is 224 MiB; scratch reuse avoids repeated allocation of that buffer, which is an allocation-size calculation rather than a measured peak-memory saving.git diff --checkand changed-Python AST checks passed. All 15 AITER slices replay to the original source tree exactly.Tests ran on 8× AMD Instinct MI355X (
gfx950, 256 CUs/device), PyTorch2.11.0+rocm7.2, ROCm 7.2.4, on 2026-10-07. GPU operator checks use AITERd78f1d4751fae6733400623b4a3ebcfe8cea6b89and SGLang50886800373fbf09fb33bd79212510541868063dunless a different checkout is explicitly named. These are combined-tree checks, not isolated PR-head A/B runs.Evidence artifacts in the benchmark bundle (
upstream/k3-quark-opts/evidence/):logs/aiter-contracts.log,direct-gpu.json,logs/direct-gpu.log,logs/direct-gpu-followup.log. These paths identify retained local raw data; they are not public artifact URLs.Submission Checklist