From 1133271808c0c64bfce40df36147dcff6490ee6c Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:55:57 +0800 Subject: [PATCH 01/12] [Ascend] Add manual and persistent staged FlashMLA examples --- examples/ascend/flashmla/LICENSE.flashmla | 21 + examples/ascend/flashmla/README.md | 178 + examples/ascend/flashmla/RESULTS.md | 340 ++ examples/ascend/flashmla/WORKTREES.md | 63 + examples/ascend/flashmla/core.py | 743 ++++ examples/ascend/flashmla/decode.py | 103 + examples/ascend/flashmla/manual.py | 26 + examples/ascend/flashmla/manual_intrinsics.h | 64 + examples/ascend/flashmla/quantization.py | 64 + .../results/benchmark-input-metadata.json | 675 ++++ .../flashmla/results/decode-bench-final.json | 1088 ++++++ .../results/decode-correctness-final.json | 116 + .../results/decode-shared-fp4-recheck.json | 168 + .../ascend/flashmla/results/environment.json | 51 + .../flashmla/results/prefill-bench-final.json | 362 ++ .../flashmla/results/prefill-code-size.json | 140 + .../results/prefill-correctness-final.json | 330 ++ .../results/staged-b4-layouts-2026-10-20.json | 343 ++ .../staged-decode-bench-2026-10-18.json | 1589 +++++++++ ...ecode-compiler-equivalence-2026-10-18.json | 633 ++++ .../staged-main-rebase-2026-10-20.json | 275 ++ ...d-main-rebase-decode-bench-2026-10-20.json | 3079 +++++++++++++++++ ...-main-rebase-prefill-bench-2026-10-20.json | 947 +++++ .../staged-manual-order-2026-10-20.json | 309 ++ examples/ascend/flashmla/run.py | 195 ++ examples/ascend/flashmla/run_decode.py | 198 ++ examples/ascend/flashmla/staged.py | 757 ++++ examples/ascend/flashmla/test_staged.py | 240 ++ 28 files changed, 13097 insertions(+) create mode 100644 examples/ascend/flashmla/LICENSE.flashmla create mode 100644 examples/ascend/flashmla/README.md create mode 100644 examples/ascend/flashmla/RESULTS.md create mode 100644 examples/ascend/flashmla/WORKTREES.md create mode 100644 examples/ascend/flashmla/core.py create mode 100644 examples/ascend/flashmla/decode.py create mode 100644 examples/ascend/flashmla/manual.py create mode 100644 examples/ascend/flashmla/manual_intrinsics.h create mode 100644 examples/ascend/flashmla/quantization.py create mode 100644 examples/ascend/flashmla/results/benchmark-input-metadata.json create mode 100644 examples/ascend/flashmla/results/decode-bench-final.json create mode 100644 examples/ascend/flashmla/results/decode-correctness-final.json create mode 100644 examples/ascend/flashmla/results/decode-shared-fp4-recheck.json create mode 100644 examples/ascend/flashmla/results/environment.json create mode 100644 examples/ascend/flashmla/results/prefill-bench-final.json create mode 100644 examples/ascend/flashmla/results/prefill-code-size.json create mode 100644 examples/ascend/flashmla/results/prefill-correctness-final.json create mode 100644 examples/ascend/flashmla/results/staged-b4-layouts-2026-10-20.json create mode 100644 examples/ascend/flashmla/results/staged-decode-bench-2026-10-18.json create mode 100644 examples/ascend/flashmla/results/staged-decode-compiler-equivalence-2026-10-18.json create mode 100644 examples/ascend/flashmla/results/staged-main-rebase-2026-10-20.json create mode 100644 examples/ascend/flashmla/results/staged-main-rebase-decode-bench-2026-10-20.json create mode 100644 examples/ascend/flashmla/results/staged-main-rebase-prefill-bench-2026-10-20.json create mode 100644 examples/ascend/flashmla/results/staged-manual-order-2026-10-20.json create mode 100644 examples/ascend/flashmla/run.py create mode 100644 examples/ascend/flashmla/run_decode.py create mode 100644 examples/ascend/flashmla/staged.py create mode 100644 examples/ascend/flashmla/test_staged.py diff --git a/examples/ascend/flashmla/LICENSE.flashmla b/examples/ascend/flashmla/LICENSE.flashmla new file mode 100644 index 0000000000..5c48bdc9f6 --- /dev/null +++ b/examples/ascend/flashmla/LICENSE.flashmla @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/examples/ascend/flashmla/README.md b/examples/ascend/flashmla/README.md new file mode 100644 index 0000000000..e1976c0d7d --- /dev/null +++ b/examples/ascend/flashmla/README.md @@ -0,0 +1,178 @@ +# FlashMLA PR #229: manual and compiler-managed Ascend scheduling + +This example reproduces the Ascend 950 sparse attention kernel in +[FlashMLA PR #229](https://github.com/deepseek-ai/FlashMLA/pull/229), pinned to +`a123d0b0191e0da7aa1e044f8644e0989cc63220`, with explicit TileLang scheduling. +The code and scheduling algorithms adapted from FlashMLA retain its MIT notice +in `LICENSE.flashmla`. + +`core.py` owns the complete AIC/AIV pipeline, online softmax, gather2 index +processing, skip-scale accumulation, and synchronization. `quantization.py` +implements FP8 / FP4 dequantization in TileLang SIMD. `manual_intrinsics.h` +only wraps individual CANN transfer / MMA instructions and SS-buffer access; +it contains no attention loop and does not call the reference kernel. +`TL_ENABLE_AUTO_SCHEDULE=False` and explicit `T.Cube()` / `T.Vector()` are used. + +`staged.py` expresses the same attention algorithm as logical tiles with +`T.Stage` and compiler-managed versions and synchronization. A persistent KV +pipeline spans query jobs. Source order and manual pipe-order constraints +preserve issue order without `T.IssuePhase` annotations. Q uses two +explicit query slots when the periodic lifetime fits the event budget +(`4 <= nb <= 15`); other shapes retain the single-query-slot path. The two +32-row KV halves have explicit L1 offsets in separate bank groups. + +QK and PV use separate four-version L0B buffers with MajorK and MajorMN +layouts. `T.assume_conflict(bq, bv)` lets the compiler give this compatible pair +one storage-ring clock and a shared 64 KiB allocation. Their copies and matrix +products use `T.copy` and `T.gemm`. PV dispatches its nonzero and zero unit flags +through two GEMM branches in one `T.Task`; the two-KV-tile path keeps the flag in +a scalar variable so rescale synchronization survives constant folding. + +For decode with FP8 main KV and FP4 extra KV, each conditional pair-copy is one +`T.Task`, and the conditional dequantization is another. This gives the two +format alternatives a common producer/consumer ownership boundary. `raw8` and +`raw4` retain separate storage, layouts, and three-slot histories: a zero-burst +copy must preserve the previous contents of the selected format's own slot. +They are not interchangeable views of one raw buffer. `dequantize(..., slot=None)` +serves these logical raw tiles; the manual kernel's explicit-slot path is unchanged. + +The manual pair gathers use two-dimensional `T.make_tensor_from_addr` views; +the staged kernel uses `T.StridedTensor`. Both use `T.copy`, +with runtime source strides and a runtime row count of zero or two. A zero +count preserves the UB contents; it does not clear them. The local +`T.assume(n_burst > 0)` is a temporary code-generation workaround that removes +the copy's positive-count guard in the current compiler while retaining the +runtime zero-or-two DMA argument. It is **not a valid input invariant**: fully +invalid pairs still have zero bursts. Recheck the generated code and invalid-pair +cases when upgrading the compiler; this needs an explicit zero-burst copy +contract to become a supported use of `T.copy`. + +The manual softmax handoff uses the existing `tl.ascend_copy_ubuf_to_cbuf` intrinsic through +`T.call_intrin`: one DMA with four 1 KiB bursts and 1 KiB source/destination gaps. +This preserves the in-place BF16 packing, shared L1 tile layout, instruction +granularity, and manual synchronization. Ordinary `T.copy` still requires +contiguous regions on this raw UB-to-L1 path. The staged kernel instead uses a +compact 32-row S tile with an annotated layout and ordinary `T.copy`. + +Supported shapes are H=64, Dqk=Dv=512, BF16 Q and output, and topk multiples of +64. Prefill takes contiguous BF16 KV and allows strided Q / indices. Decode +accepts the PR's compact FP8 main cache and optional FP8 or FP4 extra cache. +Both support attention sinks, invalid indices, and per-query topk lengths. +Packed decode pages must be contiguous, and each cache must be smaller than +4 GiB because gather uses 32-bit byte offsets. Runtime topk lengths mask the +fixed loop extent; they do not shorten this example's pipeline loop. + +## Reproduce + +This example is proposed separately from its compiler prerequisites, listed in +[WORKTREES.md](WORKTREES.md). The example branch starts at main `d85a70a`; +apply the prerequisite compiler changes before compiling the kernels. +The historical integrated compiler was based on main `876a59e`; its +[rebase validation record](results/staged-main-rebase-2026-10-20.json) covers +the initial native build, host tests, and kernel source generation. The +[subsequent full performance sweep](RESULTS.md#2026-10-20-complete-performance-sweep-after-main-rebase) +records three BF16 prefill and ten decode cases on that integrated compiler, +including output checks and same-input reference/manual/staged comparisons. + +Build TileLang with the prerequisite changes and Ascend enabled. The original performance +measurements in `RESULTS.md` used `a35f8ddf45eba16c21211ec8822d56ce5363036f` +and the earlier C-adapter pair gathers. The 2026-10-18 section reports the +native-copy manual and compiler-managed decode comparison with its exact +compiler and source identities. Build the pinned FlashMLA reference separately: + +```bash +cd /path/to/FlashMLA +MAX_JOBS=16 FLASH_MLA_BUILD_TARGET_PLATFORM=ASCEND python setup.py build_ext --inplace +``` + +From the TileLang worktree, select an idle physical NPU and run: + +```bash +export ASCEND_RT_VISIBLE_DEVICES=5 # select a free device on your machine +export PYTHONPATH=. +python examples/ascend/flashmla/run.py --reference /path/to/FlashMLA \ + --mode correctness --output results/prefill-correctness.json +python examples/ascend/flashmla/run.py --reference /path/to/FlashMLA \ + --mode bench --rounds 3 --repeat 10 --output results/prefill-bench.json +python examples/ascend/flashmla/run_decode.py --reference /path/to/FlashMLA \ + --mode correctness --output results/decode-correctness.json +python examples/ascend/flashmla/run_decode.py --reference /path/to/FlashMLA \ + --mode bench --rounds 3 --repeat 10 --output results/decode-bench.json +``` + +Both runners accept `--implementation reference`, `manual`, or `staged` to +select one implementation. The default `both` preserves the reference/manual +comparison; `all` includes the staged kernel. For example, add +`--implementation all` to the correctness or benchmark commands above to compare +all three on the same inputs. A selected implementation run alone still receives +the mathematical-reference correctness check before timing. + +The scripts use the pinned PR's input generator, numerical reference, +quantization, FLOP accounting, and `kernelkit.bench` profiler. All selected variants +receive the same tensors, including the reference's non-contiguous strides. +Each timed launch follows an 8 GB L2 flush; JIT compilation and warm-up are +excluded. They reverse implementation order on alternate rounds, require exactly 10 +attention-kernel samples per round, and report the median of three round means. +JSON files retain every launch duration. Timing covers the attention kernel; +host allocations, input reshapes, and ancillary device kernels are excluded, +as in the PR's attention-kernel performance metric. + +`--nk 8192 --detail` selects one prefill performance case and additionally +collects AIC/AIV pipe statistics. `--case N` selects a zero-based case in either runner. +The full performance sets contain three prefill cases and ten decode cases, +including the two prefill-as-decode cases. + +## Scheduling + +The manual kernel preserves these reference choices: four QK/PV tiles and sixteen gather2 DMA +instructions are unrolled; the BF16 ND-to-NZ row micro-loop is unrolled 32 +ways with post-increment pointers; softmax row groups and the persistent +pipeline stay rolled. Softmax exp/sum processes eight rows per rolled +iteration, and the output epilogue processes two rows per rolled iteration. +First-block softmax and output-accumulation conditions select specialized VFs +outside SIMD row loops, matching the reference's C++ template specialization. +The vendor flags include `-Ofast -mllvm -enable-hiipu-vf-loop-unroll`. + +The manual kernel's synchronization uses explicit buffer locks, cross-core flags, unit flags, +and the shared scalar buffer. The skip-scale notification is issued before +gather only at an index-group boundary, and after gather otherwise. Changing +this ordering measurably reduced pipeline overlap in the initial version. +For complete 8/10-block prefill index groups, two live index versions suffice; +shorter groups keep four. This keeps dynamic UB within CANN's 248 KiB budget. +Bounds legalization is disabled because index preprocessing and zero-burst +DMA implement the masks explicitly; invalid pairs must not issue a GM read. + +The staged kernel keeps QK and PV as nested scheduled loops. In the common +10-tile query, metadata leads pair gather by one iteration, and gather leads +ND-to-NZ or dequantization by two. The three KV/raw slots cover gather through +transpose/dequantization. For BF16 queries with 5–10 tiles, gather and softmax +read two query slots of index masks directly. The final softmax read precedes +reuse of its query slot; the compiler tracks the full `2 * nb` period. Other +shapes snapshot tile masks into six versions when metadata leads gather, or +five versions otherwise. The compiler emits synchronization and buffer-version +counters for these ownership intervals. Queries outside the Q-slot event +budget use the single-query-slot schedule. + +The pipelined source orders conversion, QK, softmax/statistics, metadata, +PV/output, gather, and SS publication by their intended issue order; `T.Stage` +still identifies their logical KV tile. Manual pipe-order constraints keep +gather before the scalar wait for softmax's `delta`, including when the +higher-stage SS store would otherwise tie with gather. No `T.IssuePhase` +annotations are needed. With `pipeline_stages=False`, metadata and gather +remain before their consumers, and SS is published immediately after softmax. + +`test_staged.py` covers cross-query rescaling, accumulation, partial index groups, +the flat execution path, and Q-prefetch boundaries at 5, 15, and 16 tiles. +Its independent main/extra decode length tests additionally require +`FLASHMLA_REFERENCE=/path/to/FlashMLA`. + +See `RESULTS.md` for the measured results, environment, and limitations. +The historical manual-kernel measurements and newer staged measurements are +reported separately, including the pre-rebase `9c46d1c` manual-order validation in +[staged-manual-order-2026-10-20.json](results/staged-manual-order-2026-10-20.json). +The current compiler's complete performance results are in +[prefill](results/staged-main-rebase-prefill-bench-2026-10-20.json) and +[decode](results/staged-main-rebase-decode-bench-2026-10-20.json), with every +measured launch retained. +Use the per-case runner JSON to compare correctness and +latency for the exact implementation, shape, and compiler build. diff --git a/examples/ascend/flashmla/RESULTS.md b/examples/ascend/flashmla/RESULTS.md new file mode 100644 index 0000000000..3c9571b676 --- /dev/null +++ b/examples/ascend/flashmla/RESULTS.md @@ -0,0 +1,340 @@ +# Measurement report + +Measured on 2026-10-09 with one Ascend950DT_9581 (32 AIC + 64 AIV), physical +device 5, CANN 9.2.0, Bisheng clang 15.0.5 (2026-08-27), PyTorch 2.13.0+cpu, +and torch_npu 2.13.0rc1. TileLang base: `a35f8ddf45eba16c21211ec8822d56ce5363036f`. +FlashMLA reference: `a123d0b0191e0da7aa1e044f8644e0989cc63220`. +The unchanged reference extension and manual kernel ran on the same NPU. + +## Performance + +Each entry is the median of three round means, ten launches per round, using +the PR's msprof harness and 8 GB L2 flush. A/B order alternates by round. +Speedup is reference latency divided by manual latency. Small differences +near 1.00x should be interpreted together with the recorded round variation. +These are attention-kernel times; input reshapes, allocations and other kernels +are excluded. Decode applies the same per-call flattening as the PR wrapper, +so copies of strided inputs affect the initial cache state in both variants. + +All prefill cases use nq=4096, H=64, D=512, topk=640 and attention sinks. + +| KV length | Reference μs | Manual μs | Reference TFLOPS | Manual TFLOPS | Speedup | +| --- | ---: | ---: | ---: | ---: | ---: | +| 4096 | 858.04 | 848.67 | 400.44 | 404.86 | 1.011x | +| 8192 | 853.78 | 851.15 | 402.44 | 403.68 | 1.003x | +| 32768 | 862.15 | 862.05 | 398.54 | 398.58 | 1.000x | + +Decode uses FP8 main KV, topk=128, extra topk=512, H=64, D=512 and attention +sinks. The first eight cases have sq=4, variable cache sequence lengths, +main sequence length 256 and extra sequence length 2048. The last two are the +PR's shared-cache prefill-as-decode cases, with b=1, sq=4096 and both sequence +lengths 4096. They have no topk-length tensors. + +| Extra cache | Batch | Queries/batch | Reference μs | Manual μs | Reference TFLOPS | Manual TFLOPS | Speedup | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| FP8 | 64 | 4 | 78.44 | 76.70 | 273.78 | 279.99 | 1.023x | +| FP8 | 128 | 4 | 144.46 | 142.34 | 297.32 | 301.75 | 1.015x | +| FP8 | 256 | 4 | 274.83 | 271.30 | 312.56 | 316.63 | 1.013x | +| FP8 | 512 | 4 | 539.45 | 533.56 | 318.47 | 321.99 | 1.011x | +| FP4 | 64 | 4 | 70.60 | 71.38 | 304.19 | 300.84 | 0.989x | +| FP4 | 128 | 4 | 131.34 | 131.99 | 327.00 | 325.40 | 0.995x | +| FP4 | 256 | 4 | 247.43 | 248.10 | 347.17 | 346.23 | 0.997x | +| FP4 | 512 | 4 | 489.66 | 497.20 | 350.85 | 345.53 | 0.985x | +| FP8 | 1 | 4096 | 1045.89 | 1032.69 | 328.52 | 332.72 | 1.013x | +| FP4 | 1 | 4096 | 936.84 | 953.41 | 366.76 | 360.39 | 0.983x | + +The FP4 shared-cache case had one anomalous round in the full sweep: its +manual mean was 6749.84 μs, with eight launches between 6.25 and 10.57 ms. +The other two round means were 953.41 and 953.17 μs. All samples remain in the +JSON and the table follows the same median-of-round-means rule as every case. +A separate five-round recheck gave manual round means of +952.22, 953.05, 952.43, 953.13, 952.65 μs, with median +952.65 μs, versus 939.97 μs +for the reference. The cause of the anomalous full-sweep capture is unresolved; +see [the full recheck](results/decode-shared-fp4-recheck.json). + +## Correctness and scope + +The correctness runner checks both variants against the PR's PyTorch +reference, using its output cosine/error tolerances and logits/LSE checks. +Prefill includes 14 cases: sinks, strided inputs, per-query lengths including +zero, invalid indices, partial index groups, and amplified KV values that +trigger rescaling. Decode includes 9 cases: FP8-only, FP8+FP8, FP8+FP4, +zero/all-invalid requests, per-query main/extra lengths, and odd page sizes. +Performance cases additionally compare the two implementations' outputs. + +The supported example scope is H=64, D=512 and compact KV pages. Each quantized +cache must fit in 32-bit byte offsets. Runtime topk lengths are implemented as +masks over the static topk loop, while the reference shortens the loop. Thus +variable-topk performance is not claimed by these fixed-topk measurements. + +## Controlled changes + +The diagnostic case uses nq=4096, nk=8192, topk=640. Intermediate measurements +below use one ten-launch round, except the initial measurement (three rounds). +They guide tuning; the final comparison above uses three rounds for every case. + +| Manual implementation | Latency μs | +| --- | ---: | +| Initial explicit pipeline | 1625.56 | +| Match gather / skip-scale notification ordering | 1251.82 | +| Match softmax row grouping and packed conversion | 1242.92 | +| Rely on explicit masks and zero-burst DMA | 1206.52 | +| Hoist scalar indices out of VF | 1206.47 | +| Specialize first-block / accumulation VFs and process two output rows together | 850.94 | + +The last experiment combines VF specialization with the reference's paired-row +output epilogue; it does not isolate the contribution of either change alone. +Dynamic-UB alignment had no material effect in the measured case. + +## Unroll and code size + +The source explicitly keeps persistent and SIMD row-group loops rolled, while +unrolling the same small QK/PV, gather and ND-to-NZ groups as the reference. +The native flags include `-Ofast -mllvm -enable-hiipu-vf-loop-unroll`. +The ND-to-NZ VF is 536 bytes in both binaries. For the representative prefill +specialization, total device `.text` is 14,440 bytes for the reference and +15,560 bytes for the manual version. All manual SIMD VFs report zero stack +usage; the largest reported vector-register count is 12. This inspection did +not collect hardware ICache-miss counters. + +## Evidence and commands + +[README.md](README.md) contains the exact reproduction commands. Raw launch +samples and round means are saved in [results/prefill-bench-final.json](results/prefill-bench-final.json) +and [results/decode-bench-final.json](results/decode-bench-final.json). The +correctness JSON files, source hashes, compiler version and function sizes +are in [results](results). + +These hardware results were collected on the pinned TileLang revision above. +The publication branch targets a newer `main`. The current pair gathers use +native `T.copy` with the temporary `T.assume` workaround documented in `README.md`; +the performance numbers above predate this change. The 2026-10-18 section below +reports the newer native-copy and staged decode measurements separately. +`results/environment.json` retains the original measurement and publication +source hashes as historical evidence. The current attention core and intrinsic +header differ from those snapshots. + +On 2026-10-14, the native pair-copy version passed all 14 prefill and 9 decode +correctness cases through the TileLang runtime, including zero/all-invalid +pairs. This recheck used the unchanged compiler from `b683a26`, CANN 9.2.0, +torch_npu 2.13.0rc1, and Ascend950DT_9581 physical devices 7 (prefill) and 6 +(decode). Generated code retained runtime 0/2 DMA counts and decode address +clamping. For representative BF16 prefill and FP8+FP4 decode specializations, +Bisheng assembly retained the same DMA instruction counts without increasing +conditional-jump counts. This was a correctness/code-generation recheck, +not a performance measurement. + +The full local debug evidence is under `.artifacts/flashmla-pr229/` in the +worktree: `final-suite.log`, generated `.asc` files, intermediate experiments, +`prefill-specialized-resources.log`, `prefill-specialized-symbols.txt`, +`reference-prefill-symbols.txt`, and the pinned reference checkout. These large +scratch artifacts are ignored by Git. Transient device initialization +failures were retried before the final suite; failed starts provide no samples. + +## Compiler-managed decode, 2026-10-18 + +The native-copy manual kernel and compiler-managed `staged.py` were measured +against the same pinned FlashMLA reference on the same Ascend950DT_9581, +physical device 5. Each value is the median of three round means with ten +launches per round and an 8 GB L2 flush before every launch. All selected +implementations pass the numerical checks before timing. The shape definitions +and fixed-topk scope above also apply here. + +| Extra cache | Batch | Queries/batch | Reference μs | Manual μs | Staged μs | Manual / staged | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| FP8 | 64 | 4 | 78.49 | 76.56 | 71.98 | 1.064x | +| FP8 | 128 | 4 | 144.43 | 143.03 | 136.07 | 1.051x | +| FP8 | 256 | 4 | 276.97 | 273.28 | 260.16 | 1.050x | +| FP8 | 512 | 4 | 546.29 | 540.71 | 511.31 | 1.057x | +| FP4 | 64 | 4 | 70.86 | 70.14 | 68.95 | 1.017x | +| FP4 | 128 | 4 | 130.75 | 131.15 | 129.01 | 1.017x | +| FP4 | 256 | 4 | 251.05 | 252.51 | 246.31 | 1.025x | +| FP4 | 512 | 4 | 497.03 | 500.63 | 479.67 | 1.044x | +| FP8 | 1 | 4096 | 1045.70 | 1034.47 | 974.18 | 1.062x | +| FP4 | 1 | 4096 | 937.63 | 954.60 | 935.75 | 1.020x | + +All ten staged decode cases are faster than the manual kernel in this sweep. +The staged implementation also passed the nine decode correctness cases and +four independent main/extra-length boundary cases. Raw samples are in +[staged-decode-bench-2026-10-18.json](results/staged-decode-bench-2026-10-18.json). + +These device measurements used compiler `0d86dd0b886624bb2e35bbfdf83e9e30cc358c98`. +After the pointer-layout fix in `34894d2ee0bf1c311e3587dd530ddb989ae8fa70`, all +23 unique decode instances were compiled again and loaded in fresh processes. +Their complete device executables, host sources, device sources, and parameter +files are byte-identical to the measured and correctness-tested versions; +[the compiler equivalence record](results/staged-decode-compiler-equivalence-2026-10-18.json) +contains the per-case hashes. This is a compilation equivalence check, not a +second device timing run. Kernel source and native-library hashes remain in +the benchmark provenance. + +## 2026-10-20: high-level QK/PV with a shared L0B ring + +Commit `e924481` changed `staged.py` to use independent MajorK and MajorMN B4 buffers, ordinary +`T.copy` / `T.gemm`, and default `T.assume_conflict(bq, bv)`. Compiler commit +`c536475` gives the compatible pair one ring clock while preserving its layouts. +The allocation remains 64 KiB. The original A4 and single S-to-A copy are retained. +PV uses two GEMM branches in one `T.Task`; the two-tile path holds the unit flag +in a scalar variable so intermediate rescale drains keep their M/FIX synchronization. + +The preceding best kernel is preserved in local commit `9e691b0` and in the +frozen source identified by SHA256 `1cd8a2282ab2d960623fb85c4777c0b4d1679d158b9b6c87170e8fbafde55697`. +Both versions below were built with the same compiler and compared on identical +inputs, alternating order over three rounds of ten launches on physical NPU 5. +Times are the median of the three round means; a positive change means slower. + +| BF16 prefill, nk | Previous best, μs | High-level B4, μs | Change | +| ---: | ---: | ---: | ---: | +| 4096 | 865.48 | 863.65 | -0.21% | +| 8192 | 864.46 | 866.13 | +0.19% | +| 32768 | 873.02 | 873.40 | +0.04% | + +Validation passed all 14 BF16 correctness cases, all nine decode correctness +cases, six permanent BF16 regressions, and four independent decode-length +boundary regressions. The compiler change passed 741 host tests and its device +regression with unequal owner lengths and repeated launches. This change's +performance sweep covers BF16; the earlier ten-case decode timing table above +belongs to its separately identified baseline. + +Source/native identities, correctness counts, and all timing samples are in +[staged-b4-layouts-2026-10-20.json](results/staged-b4-layouts-2026-10-20.json). +[WORKTREES.md](WORKTREES.md) records the separately proposed compiler +prerequisites and the identities of the historical integrated builds. + +## 2026-10-20: manual issue order without IssuePhase, before main rebase + +Commit `9c46d1c` removes all `T.IssuePhase` calls from FlashMLA. The source +orders operations by their intended issue order while `T.Stage` identifies +their logical KV tiles. The compiler now separates physical issue phases +when the larger-stage-first reconstruction rule would reverse a manual +source-order edge on a shared core/pipe. Generated AIV code issues gather +before waiting for softmax's delta and publishing SS. That historical compiler +still contained the unused IssuePhase API; the split compiler proposals omit it. + +Existing validation records report 68 focused tests, 776 AutoSchedule tests, +eight schedule comparisons, and 11 FlashMLA device regressions passing. +The device suite covers pipeline and flat execution, cross-query rescaling, +alternating rescale/accumulate paths, index-group tails, Q-prefetch boundaries +at 5/15/16 tiles, and four independent main/extra decode-length cases. +The source and native library checked before rebasing match these recorded +SHA256 values. The kernel source is unchanged by the rebase; the rebuilt +compiler is a separate validation target. + +```text +staged.py: 986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2 +libtilelang.so: f63a9bebad45fe496bb7e89c621d29efaf819fca41a61e612648175c65885713 +``` + +The performance comparison used Ascend950DT_9581 physical NPU 6, identical +inputs, kernelkit/msprof with L2 flushing, and three alternating rounds of +ten launches. Values below are medians of round means. Timing started after +the device regression suite completed. + +| BF16 prefill (nq, nk, topk) | Intermediate baseline, μs | Before rebase, μs | Change | +| --- | ---: | ---: | ---: | +| (4096, 4096, 640) | 866.10 | 867.44 | +0.154% | + +This baseline is the intermediate source with one remaining IssuePhase, +SHA256 `b837367b6c3a0e9d8d9faabaac3946726361012ba6bf1b87050163a5b827425e`. +Its frozen ASC from the previous compiler was recompiled after checking the +kernel ABI and returned source identity. Outputs match that baseline bitwise. +This is a separate comparison from the three-case B4 table above; it does +not directly compare against the original `e924481` or `9e691b0` source. + +[The manual-order record](results/staged-manual-order-2026-10-20.json) preserves +source/library identities, schedule comparisons, and the three round means +for each variant. Per-launch samples were not saved in that record. +The original local summary was written before committing, with `e924481` as +its parent HEAD; the validated changes were subsequently committed as `9c46d1c`. +This pre-rebase record does not include a rerun of the full 14-case BF16 or nine-case +decode correctness runners, nor the full three-case BF16 and ten-case decode +timing matrices. Earlier measurements retain their original revision scope. + +## 2026-10-20: main rebase, build and host validation + +The integrated kernel branch at this measurement was `00e7b9b`, based on upstream main `876a59e`. +The rebase removes upstreamed SS, independent-AIV, partial-copy, and SIMD/PTO +changes from the local delta. The composite prerequisite commit `35a66f4` +becomes `d3042c7`, reduced from 77 files to 22 while retaining the remaining +NZ pitch, compact-allocation, and access-footprint fixes. Main's newer codegen +base is preserved. All committed FlashMLA files are unchanged from `9c46d1c`. + +The incremental native build passed, followed by 987 host tests covering +AutoSchedule and the relevant copy, ND2NZ, L1-offset, layout, and PTO paths. +Full TileLang lowering and ASC source generation passed for BF16 +`(nq=65, nk=2048, topk=640)` and FP8 main plus FP4 extra +`(nq=65, nk=2048, topk=128, extra_topk=512)`, with variable lengths enabled. +The initial host selection included unmarked device tests, which failed at +NPU initialization with device visibility disabled. The final run explicitly +selects host tests from those layout files; both logs are retained. + +The rebuilt `libtilelang.so` SHA256 is +`489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b`. +The [rebase validation record](results/staged-main-rebase-2026-10-20.json) +contains the commit mapping, exact test command, dependency identities, +library hashes, generated-source hashes, and lowering warnings. +That initial check did not include Bisheng compilation of these generated +kernels, NPU correctness runs, or performance measurements. The subsequent +device performance sweep on the same compiler is recorded below. + +## 2026-10-20: complete performance sweep after main rebase + +The integrated `00e7b9b` branch, based on main `876a59e`, was measured on +Ascend950DT_9581 physical NPU 6 with CANN 9.2.0, Bisheng 15.0.5, +torch 2.13.0+cpu, and torch_npu 2.13.0rc1. The native library SHA256 remains +`489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b`. +All 26 manual/staged kernels were freshly compiled and passed allocation +checks before device execution. Runtime cache hits were checked against their +source and binary identities, with new compilation forbidden during the sweep. +The pinned FlashMLA reference remains `a123d0b0191e0da7aa1e044f8644e0989cc63220`. + +All three implementations used the same inputs at each of the three BF16 +prefill and ten decode shapes. Manual and staged outputs passed the existing +runners' comparisons against FlashMLA before timing. Each measurement used +kernelkit/msprof, an 8 GB buffer to flush L2, ten warmup launches, and ten measured +launches per round. Implementation order reversed on alternate rounds. +Tables report the median of three round means in microseconds; a positive +change means staged is slower than manual. All 1,170 measured attention-kernel +durations are retained. Compilation, allocations, and wrapper kernels are +excluded from the reported attention-kernel time. + +BF16 prefill uses `nq=4096, topk=640, H=64, Dqk=Dv=512`: + +| nk | Reference, μs | Manual, μs | Staged, μs | Staged vs manual | +| ---: | ---: | ---: | ---: | ---: | +| 4096 | 853.32 | 848.60 | 864.24 | +1.84% | +| 8192 | 854.87 | 850.44 | 866.68 | +1.91% | +| 32768 | 861.95 | 861.72 | 873.97 | +1.42% | + +Decode uses FP8 main KV, `topk=128`, and `extra_topk=512`. The extra-cache +format and batch/query dimensions vary: + +| Extra format | b | sq | Reference, μs | Manual, μs | Staged, μs | Staged vs manual | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| FP8 | 64 | 4 | 78.77 | 77.13 | 72.84 | -5.56% | +| FP8 | 128 | 4 | 145.32 | 143.10 | 137.16 | -4.15% | +| FP8 | 256 | 4 | 275.85 | 271.63 | 264.53 | -2.62% | +| FP8 | 512 | 4 | 541.40 | 535.04 | 519.54 | -2.90% | +| FP4 | 64 | 4 | 70.62 | 70.26 | 67.93 | -3.33% | +| FP4 | 128 | 4 | 129.65 | 130.13 | 127.20 | -2.25% | +| FP4 | 256 | 4 | 247.19 | 249.05 | 245.40 | -1.46% | +| FP4 | 512 | 4 | 490.12 | 495.86 | 479.41 | -3.32% | +| FP8 | 1 | 4096 | 1046.16 | 1034.05 | 998.09 | -3.48% | +| FP4 | 1 | 4096 | 936.10 | 953.46 | 957.61 | +0.44% | + +Staged remains 1.42–1.91% slower than manual on BF16 prefill. It is +1.46–5.56% faster on nine decode cases; FP8+FP4 with `b=1, sq=4096` is +0.44% slower than manual and 2.30% slower than the reference. +The [prefill record](results/staged-main-rebase-prefill-bench-2026-10-20.json) +and [decode record](results/staged-main-rebase-decode-bench-2026-10-20.json) +contain every launch duration, exact compile arguments, resource usage, +and source/native/generated-binary hashes. + +This sweep validates outputs at the benchmark shapes; the separate 14-case +BF16 and nine-case decode correctness runners were not rerun. Earlier tables +retain their original source/compiler/device scope. The simultaneous +reference/manual/staged measurements above do not by themselves isolate the +performance effect of rebasing. Profiler logs reported unsupported +`AICORE_SHAPE` metadata, but every capture contained the expected ten +attention-kernel records. diff --git a/examples/ascend/flashmla/WORKTREES.md b/examples/ascend/flashmla/WORKTREES.md new file mode 100644 index 0000000000..90b15da16c --- /dev/null +++ b/examples/ascend/flashmla/WORKTREES.md @@ -0,0 +1,63 @@ +# FlashMLA compiler prerequisites + +The example branch contains the kernels, runners, and historical measurement +records. Compiler changes are proposed separately against `Denverjin/tilelang:main`, +synchronized to upstream `d85a70a97218fec25856f8b5c969f167ad96dc1e` on 2026-10-21. +Apply the compiler prerequisites before building the staged kernel. +The source does not use `T.IssuePhase`, and its API is excluded from these proposals. + +| Compiler change | Fork branch | +| --- | --- | +| Support native NZ output for L0C-to-UB copies | [codex/ascend-l0c-ub-nz](https://github.com/Denverjin/tilelang/tree/codex/ascend-l0c-ub-nz) | +| Preserve Cube configuration phases across manual stages | [codex/ascend-cube-register-phases](https://github.com/Denverjin/tilelang/tree/codex/ascend-cube-register-phases) | +| Select AIV streams outside scheduled pipelines | [codex/ascend-aiv-stream-selection](https://github.com/Denverjin/tilelang/tree/codex/ascend-aiv-stream-selection) | +| Preserve complete periodic buffer lifetimes | [codex/ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/codex/ascend-periodic-buffer-lifetimes) | +| Establish buffer placement hints from exact L1 offsets | [codex/ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/codex/ascend-buffer-placement-hints) | +| Preserve physical views and pointer units during layout lowering | [codex/ascend-layout-pointer-units](https://github.com/Denverjin/tilelang/tree/codex/ascend-layout-pointer-units) | +| Reuse storage and L0B rings for explicit conflict hints | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | +| Preserve manual source order across physical phase ties | [codex/ascend-manual-pipe-order](https://github.com/Denverjin/tilelang/tree/codex/ascend-manual-pipe-order) | + +The additional NZ pitch, compact-allocation, and access-footprint prerequisite +is already proposed upstream as [tile-ai/tilelang#3530](https://github.com/tile-ai/tilelang/pull/3530). +It is not duplicated as a new fork PR. The pointer-unit branch also carries +physical allocation/view remapping fixes that are outside that upstream PR. + +The buffer-placement branch currently implements exact L1 byte offsets. +Its PR describes the intended extension to manual placement hints for every +buffer kind. That extension is not implemented yet; the kernel still calls +`T.annotate_l1_offsets` and uses the existing L1 placement contract. + +The conflict-hint changes have two parts: permission for pairwise storage reuse +and a shared version clock/ownership protocol for a compatible isolated L0B pair. +Both are needed for the separate QK/PV layouts to share a physical 64 KiB ring. +The raw FP8 and FP4 buffers retain independent storage and three-slot histories. + +## Historical integrated measurements + +The implementation source in this branch is unchanged from the integrated +kernel at `00e7b9bb64c892c5323f42bede764e8940ceb4cd`. +That compiler was based on upstream main `876a59e`; its documentation was +committed as `6bdfe6e`. The historical integration still contained an unused +IssuePhase API. These identities describe the recorded experiments, not the +compiler available on an unmodified checkout of the new example branch. + +- [RESULTS.md](RESULTS.md) contains the historical measurement tables and limitations. +- [Rebase record](results/staged-main-rebase-2026-10-20.json) records the native build, host checks, and lowering identities. +- [Prefill samples](results/staged-main-rebase-prefill-bench-2026-10-20.json) and [decode samples](results/staged-main-rebase-decode-bench-2026-10-20.json) retain the complete measured launches. +- The reference is FlashMLA commit `a123d0b0191e0da7aa1e044f8644e0989cc63220`. + +The benchmark records are not performance measurements of the independently +split compiler branches. Re-run the documented commands on a chosen integrated +compiler revision before attributing those timings to a new build. + +## Dependencies already in main + +SS buffer versioning and fixed-version UB reuse (#3514), independent AIV +execution and cross-core synchronization (#3485), partial L0 copies and layout +lookup (#3511), and the SIMD/PTO code-generation foundations are already upstream. +No additional TVM or CUTLASS modifications are required. + +The optional zero-burst API and pipeline-peeling experiments are not used. +The kernel retains its documented zero-burst code-generation workaround. +PV unit-flag selection remains in kernel code rather than requiring a separate +dynamic unit-flag compiler patch. diff --git a/examples/ascend/flashmla/core.py b/examples/ascend/flashmla/core.py new file mode 100644 index 0000000000..1350389e0d --- /dev/null +++ b/examples/ascend/flashmla/core.py @@ -0,0 +1,743 @@ +"""FlashMLA V4.1 sparse attention, explicit AIC/AIV pipeline, no AutoSchedule. + +The schedule, gather2 and skip-scale algorithm follow FlashMLA PR #229. +All loops, buffer ownership, flags, and SIMD calculations are expressed in +TileLang. manual_intrinsics.h only adapts individual CANN instructions and SS +buffer access; it contains no attention kernel or scheduling loop. +""" + +from pathlib import Path +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.layout import make_ascend_l0c_layout +from examples.ascend.flashmla.quantization import dequantize + +PASS_CONFIG = { + tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False, + tilelang.PassConfigKey.TL_DISABLE_SHARED_MEMORY_REUSE: True, + # Explicit index masks and zero-burst DMA protect every masked gather. + tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, +} + + +def rolled(n): + """Emit pragma unroll 1, matching the reference rolled loops.""" + return T.unroll(n, annotations={"pragma_unroll_factor": 1}) + + +def sparse_attention( + nq, + nk, + topk, + *, + sink=True, + variable_lengths=False, + scale=512**-0.5, + q_strides=None, + index_stride=None, + kv_format="bf16", + extra_topk=0, + extra_format="fp8", + extra_index_stride=None, + variable_extra_lengths=False, + kv_storage_bytes=1, + extra_storage_bytes=1, + cache_hint=0, +): + if topk <= 0 or topk % 64: + raise ValueError("topk must be a positive multiple of 64") + if kv_format not in ("bf16", "fp8") or extra_format not in ("fp8", "fp4"): + raise ValueError("Supported formats are BF16 prefill and FP8 main / FP8 or FP4 extra decode") + if extra_topk < 0 or extra_topk % 64: + raise ValueError("extra_topk must be a nonnegative multiple of 64") + if max(kv_storage_bytes, extra_storage_bytes) >= 2**32: + raise ValueError("The manual gather uses 32-bit byte offsets per KV cache") + decode = kv_format != "bf16" + fp4 = kv_format == "fp4" + extra_fp4 = extra_format == "fp4" + record_bytes = 288 if fp4 else 528 + extra_record_bytes = 288 if extra_fp4 else 528 + valid_limit = T.uint32(0x80000000) if decode else nk + extra_index_stride = extra_index_stride or max(extra_topk, 1) + if not decode and nk >= 2**21: + raise ValueError("This gather2 implementation uses signed 32-bit byte offsets") + q_strides = q_strides or (64 * 512, 512, 1) + index_stride = index_stride or topk + nb0 = topk // 64 + nb1 = (extra_topk + 63) // 64 + nb = nb0 + nb1 + ib = 640 if decode else (topk if topk <= 640 and topk % 128 == 0 else 128) + groups0 = (nb0 + ib // 64 - 1) // (ib // 64) + groups = groups0 + (nb1 + ib // 64 - 1) // (ib // 64) + # A full 8/10-block index group needs only two live versions. This keeps + # dynamic UB below 248 KiB, leaving CANN its kernel-reserved 8 KiB. + index_versions = 2 if not decode and ib >= 512 and nb % (ib // 64) == 0 else 4 + source = Path(__file__).with_name("manual_intrinsics.h").read_text() + + @T.macro + def cset(pipe, flag): + T.ascend_cross_core_set_flag(4, pipe, flag) + T.ascend_cross_core_set_flag(4, pipe, flag + 16) + + @T.macro + def cwait(pipe, flag): + T.ascend_cross_core_wait_flag(4, pipe, flag) + T.ascend_cross_core_wait_flag(4, pipe, flag + 16) + + @T.macro + def load_q(job, bx, Q, l1): + qidx = bx + job * 32 + T.ascend_wait_flag("MTE1_MTE2", job % 2) + T.call_extern("void", "fm_q_load", T.access_ptr(l1[(job % 2) * 16384], "w"), T.access_ptr(Q[qidx, 0, 0], "r"), q_strides[1] * 2) + T.call_extern( + "void", "fm_q_load", T.access_ptr(l1[229376 + (job % 2) * 16384], "w"), T.access_ptr(Q[qidx, 32, 0], "r"), q_strides[1] * 2 + ) + T.ascend_set_flag("MTE2_MTE1", job % 2) + + @T.macro + def qk(job, k, acounter, l1, a, b, c, p_nz): + T.call_extern("void", "asc_set_mmad_direction_m") + slot = (job * nb + k) % 4 + for d in T.unroll(4): + ai = (acounter + d) % 4 + if d == 0 and k == 0: + T.ascend_wait_flag("MTE2_MTE1", job % 2) + T.ascend_get_buf("PIPE_MTE1", ai) + T.call_extern("void", "fm_a_load", T.access_ptr(a[ai * 8192], "w"), T.access_ptr(l1[(job % 2) * 16384], "r"), d) + T.call_extern("void", "fm_a_load", T.access_ptr(a[ai * 8192 + 512], "w"), T.access_ptr(l1[229376 + (job % 2) * 16384], "r"), d) + T.ascend_rls_buf("PIPE_MTE1", ai) + if d == 3 and k == nb - 1: + T.ascend_set_flag("MTE1_MTE2", job % 2) + if d == 0: + cwait("PIPE_MTE1", 6 + slot) + T.ascend_get_buf("PIPE_MTE1", 4 + d) + T.call_extern("void", "fm_k_load", T.access_ptr(b[d * 8192], "w"), T.access_ptr(l1[32768 + slot * 16384], "r"), d) + T.call_extern("void", "fm_k_load", T.access_ptr(b[d * 8192 + 512], "w"), T.access_ptr(l1[147456 + slot * 16384], "r"), d) + T.ascend_rls_buf("PIPE_MTE1", 4 + d) + T.ascend_get_buf("PIPE_M", ai) + T.ascend_get_buf("PIPE_M", 4 + d) + T.call_extern( + "void", + "fm_mmad", + T.access_ptr(c[8, 0, 0], "rw"), + T.access_ptr(a[ai * 8192], "r"), + T.access_ptr(b[d * 8192], "r"), + 64, + 128, + 3 if d == 3 else 2, + d == 0, + ) + T.ascend_rls_buf("PIPE_M", ai) + T.ascend_rls_buf("PIPE_M", 4 + d) + cwait("PIPE_FIX", 0) + T.dual_copy(c[8, :, :], p_nz, unit_flag_ctrl=3) + cset("PIPE_FIX", 0) + + @T.macro + def copy_o(last, c, out): + if not last: + T.ascend_set_flag("M_FIX", 0) + T.ascend_wait_flag("M_FIX", 0) + for d in T.unroll(4): + cwait("PIPE_FIX", 4 + d % 2) + T.call_extern("void", "fm_o_store", T.access_ptr(out[d % 2, 0], "w"), T.access_ptr(c[2 * d, 0, 0], "r"), 3 if last else 0) + cset("PIPE_FIX", 4 + d % 2) + if not last: + T.ascend_set_flag("FIX_M", 0) + T.ascend_wait_flag("FIX_M", 0) + + @T.macro + def pv(job, k, ai, l1, a, b, c, out): + T.call_extern("void", "asc_set_mmad_direction_n") + slot = (job * nb + k) % 4 + clear = T.alloc_var("int32", init=0) + for d in T.unroll(4): + T.ascend_get_buf("PIPE_MTE1", 4 + d) + T.call_extern("void", "fm_v_load", T.access_ptr(b[d * 8192], "w"), T.access_ptr(l1[32768 + slot * 16384], "r"), d) + T.call_extern("void", "fm_v_load", T.access_ptr(b[d * 8192 + 4096], "w"), T.access_ptr(l1[147456 + slot * 16384], "r"), d) + T.ascend_rls_buf("PIPE_MTE1", 4 + d) + if d == 3: + cset("PIPE_MTE1", 6 + slot) + if d == 0: + cwait("PIPE_MTE1", 1) + T.ascend_get_buf("PIPE_MTE1", ai) + T.call_extern("void", "fm_s_load", T.access_ptr(a[ai * 8192], "w"), T.access_ptr(l1[129024], "r")) + T.ascend_rls_buf("PIPE_MTE1", ai) + cset("PIPE_MTE1", 1) + T.ascend_get_buf("PIPE_M", ai) + if k > 0: + cwait("PIPE_S", 2) + cset("PIPE_S", 2) + clear = T.call_extern("int32", "fm_ss_get", (job % 8) * 8 + k % 8) + if clear: + copy_o(False, c, out) + T.ascend_get_buf("PIPE_M", 4 + d) + T.call_extern( + "void", + "fm_mmad", + T.access_ptr(c[2 * d, 0, 0], "rw"), + T.access_ptr(a[ai * 8192], "r"), + T.access_ptr(b[d * 8192], "r"), + 128, + 64, + T.if_then_else(k == nb - 1, 3, T.if_then_else(k == 0, 2, 0)), + (k == 0) | (clear != 0), + ) + T.ascend_rls_buf("PIPE_M", 4 + d) + T.ascend_rls_buf("PIPE_M", ai) + if k == nb - 1: + copy_o(True, c, out) + + @T.macro + def softmax_vf(first, job4, k4, idx_slot, idx_off, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta): + with T.SimdVF(): + full = T.simd.pset(32) + eight = T.simd.pset(32, "PAT_VL8") + one = T.simd.pset(32, "PAT_VL1") + # Paired indices are sorted by gather2; repeat 16 columns. + masks = T.alloc_local((4,), "boolx256") + for col in T.unroll(4): + lo = T.simd.vld(smaller[idx_slot, idx_off + col * 8], "BLK") + hi = T.simd.vld(larger[idx_slot, idx_off + col * 8], "BLK") + paired, ignored = T.simd.vintlv(lo, hi) + masks[col] = T.simd.vcmps(paired, valid_limit, op="lt") + old_sum = T.simd.vld(sums[job4 % 2, 0]) + old_scale_max = T.simd.vld(scaling_max[job4 % 2, 0]) + for r in rolled(8): + rm = T.alloc_var("float32x64") + rm = T.simd.vdup(-3.4028234663852886e38, "float32") + if not first: + rm = T.simd.vld(maxima[job4 % 2, r * 8], "E2B_B32") + max_ptr = T.simd.make_ubuf_ptr(T.access_ptr(p[r * 64], "r", extent=1537), "float32") + for col in T.unroll(4): + x, max_ptr = T.simd.vld(max_ptr, post_inc=512) + rm = T.simd.vmax(rm, x, masks[col], mode="MODE_MERGING") + reduced = T.simd.vcgmax(rm) + even, odd = T.simd.vdintlv(reduced, reduced) + rmax = T.simd.vmax(even, odd) + repeated, ignored2 = T.simd.vintlv(rmax, rmax) + T.simd.vsts(maxima[job4 % 2, r * 8], repeated, eight) + T.simd.mem_bar("VST_VLD") + new_max = T.simd.vld(maxima[job4 % 2, 0]) + chosen_max = T.alloc_var("float32x64") + factor = T.alloc_var("float32x64") + chosen_max = new_max + factor = T.simd.vdup(0.0, "float32") + if not first: + diff = T.simd.vsub(new_max, old_scale_max) + greatest = T.simd.vcmax(diff) + broadcast = T.simd.vdupv(greatest) + rescale = T.simd.vcmps(broadcast, 6.0 / scale, op="gt") + chosen_max = T.simd.vsel(new_max, old_scale_max, rescale) + factor = T.simd.vexpdif(T.simd.vmuls(old_scale_max, scale), T.simd.vmuls(chosen_max, scale)) + T.simd.vsts(delta[job4 % 2, k4 % 2, 0], greatest, one) + T.simd.vsts(factors[job4 % 2, k4 % 2, 0], factor) + T.simd.vsts(scaling_max[job4 % 2, 0], chosen_max) + T.simd.mem_bar("VST_VLD") + pp = T.simd.make_ubuf_ptr(T.access_ptr(p[0], "r", extent=2048), "float32") + for r in rolled(4): + rmax2 = T.alloc_local((2,), "float32x64") + rsum2 = T.alloc_local((2,), "float32x64") + packed2 = T.alloc_local((2, 2), "bfloat16x128") + for rr in T.unroll(2): + rmax2[rr] = T.simd.vmuls(T.simd.vld(scaling_max[job4 % 2, r * 16 + rr * 8], "E2B_B32"), scale) + for pair in T.unroll(2): + for cc in T.unroll(2): + for rr in T.unroll(2): + increment = T.if_then_else(rr == 0, 64, T.if_then_else(pair == 1 and cc == 1, -1472, 448)) + x, pp = T.simd.vld(pp, post_inc=increment) + prob = T.simd.vexpdif(T.simd.vmuls(x, scale), rmax2[rr], masks[pair * 2 + cc]) + if pair == 0 and cc == 0: + rsum2[rr] = prob + else: + rsum2[rr] = T.simd.vadd(rsum2[rr], prob) + if cc == 0: + packed2[rr, cc] = T.simd.vcvt(prob, "bfloat16", full, part=0) + else: + packed2[rr, cc] = T.simd.vcvt(prob, "bfloat16", full, part=1) + merged0 = T.simd.vadd(packed2[0, 0], packed2[0, 1]) + merged1 = T.simd.vadd(packed2[1, 0], packed2[1, 1]) + col0, col1 = T.simd.vdintlv(merged0, merged1) + T.simd.vsts(s[pair * 2048 + r * 128], col0) + T.simd.vsts(s[pair * 2048 + 1024 + r * 128], col1) + for rr in T.unroll(2): + reduced_sum = T.simd.vcpadd(T.simd.vcgadd(rsum2[rr])) + repeated_sum, ignored3 = T.simd.vintlv(reduced_sum, reduced_sum) + T.simd.vsts(sums[job4 % 2, r * 16 + rr * 8], repeated_sum, eight) + T.simd.mem_bar("VST_VLD") + if not first: + current_sum = T.simd.vld(sums[job4 % 2, 0]) + updated_sum = T.simd.vadd(T.simd.vmul(old_sum, factor), current_sum) + T.simd.vsts(sums[job4 % 2, 0], updated_sum) + + @T.macro + def scale_output_vf(add, d, job5, k5, factors, out, accum): + with T.SimdVF(): + for r in rolled(32): + factor5 = T.simd.vld(factors[job5 % 2, k5 % 2, r * 2], "BRC_B32") + for j in T.unroll(2): + cur = T.alloc_var("float32x64", init=T.simd.vld(out[d % 2, r * 128 + j * 64])) + if add: + old = T.simd.vld(accum[r * 512 + d * 128 + j * 64]) + cur = T.simd.vadd(old, cur) + T.simd.vsts(accum[r * 512 + d * 128 + j * 64], T.simd.vmul(cur, factor5)) + + @T.macro + def final_output_vf(add, d, job5, denorm, out, accum, final): + with T.SimdVF(): + mask16 = T.simd.pset(16) + mask32 = T.simd.pset(32) + pp = T.simd.make_ubuf_ptr(T.access_ptr(out[d % 2, 0], "r", extent=4096), "float32") + for r in rolled(16): + packed = T.alloc_local((2, 2), "bfloat16x128") + for rr in T.unroll(2): + norm = T.simd.vld(denorm[job5 % 2, r * 2 + rr], "BRC_B32") + for cc in T.unroll(2): + x, pp = T.simd.vld(pp, post_inc=64) + if add: + old = T.simd.vld(accum[(r * 2 + rr) * 512 + d * 128 + cc * 64]) + x = T.simd.vadd(x, old, mask32) + scaled = T.simd.vmul(x, norm, mask32) + if rr == 0: + packed[rr, cc] = T.simd.vcvt(scaled, "bfloat16", mask32, part=0) + else: + packed[rr, cc] = T.simd.vcvt(scaled, "bfloat16", mask32, part=1) + merged0 = T.simd.vadd(packed[0, 0], packed[1, 0], mask16) + merged1 = T.simd.vadd(packed[0, 1], packed[1, 1], mask16) + row0, row1 = T.simd.vdintlv(merged0, merged1) + T.simd.vsts(final[r * 2048 + d * 256], row0, mask16) + T.simd.vsts(final[r * 2048 + 1024 + d * 256], row1, mask16) + + @T.prim_func + def main( + Q: T.StridedTensor((nq, 64, 512), q_strides, "bfloat16"), + KV: T.Tensor((kv_storage_bytes,) if decode else (nk, 512), "uint8" if decode else "bfloat16"), + Indices: T.StridedTensor((nq, topk), (index_stride, 1), "int32"), + Sink: T.Tensor((64,), "float32"), + Lengths: T.Tensor((nq,), "int32"), + ExtraKV: T.Tensor((extra_storage_bytes,), "uint8"), + ExtraIndices: T.StridedTensor((nq, max(extra_topk, 1)), (extra_index_stride, 1), "int32"), + ExtraLengths: T.Tensor((nq,), "int32"), + O: T.Tensor((nq, 64, 512), "bfloat16"), + Max: T.Tensor((nq, 64), "float32"), + LSE: T.Tensor((nq, 64), "float32"), + ): + with T.Kernel(min(nq, 32)) as bx: + T.import_source(source) + # Explicit physical L1 plan balances both 256 KiB bank groups. + l1 = T.alloc_l1((262144,), "bfloat16") + a = T.alloc_l0a((32768,), "bfloat16") + b = T.alloc_l0b((32768,), "bfloat16") + # Eight 64x64 tiles hold the output; the ninth holds QK scores. + c = T.alloc_l0c((9, 64, 64), "float32") + kv = T.alloc_shared((3, 32, 1 if decode else 512), "bfloat16") + raw8 = T.alloc_shared((3, 32, 544 if decode else 1), "uint8") + raw4 = T.alloc_shared((3, 32, 288 if decode else 1), "uint8") + scales = T.alloc_shared((1024 if decode else 1,), "bfloat16") + nz = T.alloc_shared((33 * 512,), "bfloat16") + accum = T.alloc_shared((32 * 512,), "float32") + out = T.alloc_shared((2, 32 * 128), "float32") + p_nz = T.alloc_shared((32, 64), "float32") + T.annotate_layout({p_nz: make_ascend_l0c_layout(p_nz)}) + # Softmax and its in-place BF16 output use physical NZ offsets. + # Keep their flat view separate from logical layout remapping. + p = T.make_tensor_from_addr(T.access_ptr(p_nz[0, 0], "rw", extent=32 * 64), (32 * 64,), "float32", storage_scope="shared.dyn") + s = tvm.tirx.decl_buffer((32 * 128,), "bfloat16", data=p.data, scope="shared.dyn") + final = tvm.tirx.decl_buffer((32 * 1024,), "bfloat16", data=accum.data, scope="shared.dyn") + maxima = T.alloc_shared((2, 64), "float32") + scaling_max = T.alloc_shared((2, 64), "float32") + sums = T.alloc_shared((2, 64), "float32") + denorm = T.alloc_shared((2, 64), "float32") + factors = T.alloc_shared((2, 2, 64), "float32") + delta = T.alloc_shared((2, 2, 64), "float32") + attn_sink = T.alloc_shared((64,), "float32") + indices = T.alloc_shared((ib,), "int32") + smaller = T.alloc_shared((index_versions, ib // 2), "uint32") + larger = T.alloc_shared((index_versions, ib // 2), "uint32") + offsets = T.alloc_shared((ib // 2,), "uint32") + strides = T.alloc_shared((ib // 2,), "uint32") + T.call_extern("void", "fm_init") + + with T.Cube(): + cset("PIPE_S", 1) + for i in T.unroll(4): + cset("PIPE_S", 6 + i) + for i in T.unroll(2): + T.ascend_set_flag("MTE1_MTE2", i) + acounter = T.alloc_var("int32", init=0) + jobs = T.ceildiv(nq - bx, 32) + load_q(0, bx, Q, l1) + qk(0, 0, acounter, l1, a, b, c, p_nz) + for job in rolled(jobs): + if job + 1 < jobs: + load_q(job + 1, bx, Q, l1) + for k in rolled(nb - 1): + qk(job, k + 1, acounter, l1, a, b, c, p_nz) + pv(job, k, acounter, l1, a, b, c, out) + acounter = (acounter + 1) % 4 + if job + 1 < jobs: + qk(job + 1, 0, acounter, l1, a, b, c, p_nz) + pv(job, nb - 1, acounter, l1, a, b, c, out) + acounter = (acounter + 1) % 4 + T.ascend_pipe_barrier("PIPE_ALL") + + with T.Vector() as sid: + if sink: + T.copy(Sink[sid * 32 : sid * 32 + 32], attn_sink[:32]) + T.ascend_set_flag("MTE2_V", 2) + T.ascend_cross_core_set_flag(4, "PIPE_S", 0) + for i in T.unroll(2): + T.ascend_cross_core_set_flag(4, "PIPE_S", 4 + i) + for i in rolled(3): + T.ascend_get_buf("PIPE_V", 8 + i) + with T.SimdVF(): + if decode: + zero = T.simd.vdup(0, "uint8") + for j in rolled(68): + T.simd.vsts(raw8[i, j * 256 // 544, j * 256 % 544], zero) + for j in rolled(36): + T.simd.vsts(raw4[i, j * 256 // 288, j * 256 % 288], zero) + else: + zero = T.simd.vdup(0.0, "bfloat16") + for j in rolled(128): + T.simd.vsts(kv[i, j // 4, j % 4 * 128], zero) + T.ascend_rls_buf("PIPE_V", 8 + i) + emitted = T.alloc_var("int32", init=0) + total = T.ceildiv(nq - bx, 32) * nb + for step in rolled(total + 5): + if not decode: # noqa: SIM102 - compile-time format dispatch + # Stage 0: fetch the next index group. + if step < total: + job0 = step // nb + k0 = step % nb + q0 = bx + job0 * 32 + extra0 = k0 >= nb0 + rel0 = T.if_then_else(extra0, k0 - nb0, k0) + if rel0 % (ib // 64) == 0: + T.ascend_get_buf("PIPE_MTE2", 0) + if extra_topk > 0 and extra0: + length0 = T.min(ib, extra_topk - rel0 * 64) + T.copy(ExtraIndices[q0, rel0 * 64 : rel0 * 64 + length0], indices[:length0]) + else: + length0 = T.min(ib, topk - rel0 * 64) + T.copy(Indices[q0, rel0 * 64 : rel0 * 64 + length0], indices[:length0]) + T.ascend_rls_buf("PIPE_MTE2", 0) + # Stage 2: ND -> compact NZ, then publish a KV slot to AIC. + if step >= 2 and step - 2 < total: + t2 = step - 2 + slot2 = T.alloc_var("int32", init=t2 % 3) + T.ascend_get_buf("PIPE_V", 8 + slot2) + T.ascend_get_buf("PIPE_V", 11) + if decode: + if extra_topk > 0 and t2 % nb >= nb0: + if extra_fp4: + dequantize(raw4, scales, nz, slot2, True) + else: + dequantize(raw8, scales, nz, slot2, False) + else: + if fp4: + dequantize(raw4, scales, nz, slot2, True) + else: + dequantize(raw8, scales, nz, slot2, False) + else: + with T.SimdVF(): + nd_mask = T.simd.pset(16) + src_ptr = T.simd.make_ubuf_ptr(T.access_ptr(kv[slot2, 0, 0], "r", extent=16384), "bfloat16") + dst0 = T.simd.make_ubuf_ptr(T.access_ptr(nz[0], "w", extent=16896), "bfloat16") + dst1 = T.simd.make_ubuf_ptr(T.access_ptr(nz[4224], "w", extent=12672), "bfloat16") + for _col in rolled(2): + for row in T.unroll(32): + x0, src_ptr = T.simd.vld(src_ptr, post_inc=128) + x1, src_ptr = T.simd.vld(src_ptr, post_inc=T.if_then_else(row == 31, -15744, 384)) + config = (33 << 16) | T.if_then_else(row == 31, 497, 1) + dst0 = T.simd.vsstb(x0, dst0, config, nd_mask, update=True) + dst1 = T.simd.vsstb(x1, dst1, config, nd_mask, update=True) + T.ascend_rls_buf("PIPE_V", 11) + T.ascend_rls_buf("PIPE_V", 8 + slot2) + T.ascend_get_buf("PIPE_MTE3", 11) + T.ascend_cross_core_wait_flag(4, "PIPE_MTE3", 6 + t2 % 4) + T.call_extern( + "void", + "fm_kv_push", + T.access_ptr(l1[T.if_then_else(sid == 0, 32768, 147456) + t2 % 4 * 16384], "w"), + T.access_ptr(nz[0], "r"), + ) + T.ascend_cross_core_set_flag(4, "PIPE_MTE3", 6 + t2 % 4) + T.ascend_rls_buf("PIPE_MTE3", 11) + # Stage 4: online softmax, with a threshold of 6 for rescaling. + if step >= 4 and step - 4 < total: + t4 = step - 4 + job4 = T.alloc_var("int32", init=t4 // nb) + k4 = T.alloc_var("int32", init=t4 % nb) + extra4 = k4 >= nb0 + rel4 = T.if_then_else(extra4, k4 - nb0, k4) + idx_slot = T.alloc_var( + "int32", init=(job4 * groups + T.if_then_else(extra4, groups0, 0) + rel4 // (ib // 64)) % index_versions + ) + idx_off = T.alloc_var("int32", init=(rel4 % (ib // 64)) * 32) + T.ascend_cross_core_wait_flag(4, "PIPE_V", 0) + T.ascend_get_buf("PIPE_V", 2 + job4 % 2) + if k4 == 0: + softmax_vf(True, job4, k4, idx_slot, idx_off, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta) + else: + softmax_vf(False, job4, k4, idx_slot, idx_off, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta) + T.ascend_rls_buf("PIPE_V", 2 + job4 % 2) + T.ascend_set_flag("V_MTE3", 0) + if k4 > 0: + T.ascend_set_flag("V_S", 0) + T.ascend_wait_flag("V_MTE3", 0) + T.ascend_cross_core_wait_flag(4, "PIPE_MTE3", 1) + # Preserve one DMA for four 1 KiB NZ halves. Length + # and gaps are in 32-byte blocks; each view spans 7 KiB. + T.call_intrin( + "void", + tvm.tirx.op.Op.get("tl.ascend_copy_ubuf_to_cbuf"), + T.access_ptr(l1[129024 + sid * 512], "w", extent=3584), + T.access_ptr(s[0], "r", extent=3584), + 0, # Reserved sub-block ID. + 4, # Bursts. + 32, # Blocks per burst. + 32, # Source gap: in-place BF16 packing. + 32, # Destination gap: the other AIV's half. + ) + T.ascend_cross_core_set_flag(4, "PIPE_MTE3", 1) + T.ascend_cross_core_set_flag(4, "PIPE_MTE3", 0) + if k4 == nb - 1: + if sink and job4 == 0: + T.ascend_wait_flag("MTE2_V", 2) + T.ascend_get_buf("PIPE_V", 2 + job4 % 2) + with T.SimdVF(): + rmax, unused = T.simd.vdintlv(T.simd.vld(maxima[job4 % 2, 0]), T.simd.vld(maxima[job4 % 2, 0])) + smax, unused1 = T.simd.vdintlv(T.simd.vld(scaling_max[job4 % 2, 0]), T.simd.vld(scaling_max[job4 % 2, 0])) + rsum, unused2 = T.simd.vdintlv(T.simd.vld(sums[job4 % 2, 0]), T.simd.vld(sums[job4 % 2, 0])) + invalid = T.simd.vcmps(rmax, -3.4028234663852886e38, op="eq") + maximum = T.simd.vmuls(rmax, scale) + scaling = T.simd.vmuls(smax, scale) + logsum = T.simd.vadd(T.simd.vln(rsum), scaling) + denom = T.alloc_var("float32x64", init=rsum) + if sink: + sinkprob = T.simd.vexpdif(T.simd.vld(attn_sink[0]), scaling) + denom = T.simd.vadd(denom, sinkprob) + reciprocal = T.simd.vdiv(T.simd.vdup(1.0, "float32"), denom) + T.simd.vsts(denorm[job4 % 2, 0], T.simd.vsel(T.simd.vdup(0.0, "float32"), reciprocal, invalid)) + T.simd.vsts(maxima[job4 % 2, 0], T.simd.vsel(T.simd.vdup(float("-inf"), "float32"), maximum, invalid)) + T.simd.vsts(sums[job4 % 2, 0], T.simd.vsel(T.simd.vdup(float("inf"), "float32"), logsum, invalid)) + T.ascend_rls_buf("PIPE_V", 2 + job4 % 2) + T.ascend_get_buf("PIPE_MTE3", 2 + job4 % 2) + if not decode: + T.copy(maxima[job4 % 2, :32], Max[bx + job4 * 32, sid * 32 : sid * 32 + 32]) + T.copy(sums[job4 % 2, :32], LSE[bx + job4 * 32, sid * 32 : sid * 32 + 32]) + T.ascend_rls_buf("PIPE_MTE3", 2 + job4 % 2) + if decode: # noqa: SIM102 - compile-time format dispatch + # Stage 0: fetch the next index group. + if step < total: + job0 = step // nb + k0 = step % nb + q0 = bx + job0 * 32 + extra0 = k0 >= nb0 + rel0 = T.if_then_else(extra0, k0 - nb0, k0) + if rel0 % (ib // 64) == 0: + T.ascend_get_buf("PIPE_MTE2", 0) + if extra_topk > 0 and extra0: + length0 = T.min(ib, extra_topk - rel0 * 64) + T.copy(ExtraIndices[q0, rel0 * 64 : rel0 * 64 + length0], indices[:length0]) + else: + length0 = T.min(ib, topk - rel0 * 64) + T.copy(Indices[q0, rel0 * 64 : rel0 * 64 + length0], indices[:length0]) + T.ascend_rls_buf("PIPE_MTE2", 0) + # Prepare gather2 parameters for the selected cache segment. + job0 = step // nb + k0 = step % nb + extra0 = k0 >= nb0 + rel0 = T.if_then_else(extra0, k0 - nb0, k0) + if step < total and rel0 % (ib // 64) == 0: + idx0 = T.alloc_var( + "int32", init=(job0 * groups + T.if_then_else(extra0, groups0, 0) + rel0 // (ib // 64)) % index_versions + ) + main_length = T.min(topk, Lengths[bx + job0 * 32]) if variable_lengths else topk + extra_length = T.min(extra_topk, ExtraLengths[bx + job0 * 32]) if variable_extra_lengths else extra_topk + valid_len = T.alloc_var("int32", init=T.max(T.if_then_else(extra0, extra_length, main_length) - rel0 * 64, 0)) + T.ascend_get_buf("PIPE_V", 0) + T.ascend_get_buf("PIPE_V", 1) + with T.SimdVF(): + for i in rolled(ib // 128): + x0 = T.reinterpret("uint32x64", T.simd.vld(indices[i * 128])) + x1 = T.reinterpret("uint32x64", T.simd.vld(indices[i * 128 + 64])) + pos0 = T.simd.vci(i * 128, "int32") + pos1 = T.simd.vci(i * 128 + 64, "int32") + sentinel = T.simd.vdup(valid_limit, "uint32") + x0 = T.simd.vsel(x0, sentinel, T.simd.vcmps(pos0, valid_len, op="lt")) + x1 = T.simd.vsel(x1, sentinel, T.simd.vcmps(pos1, valid_len, op="lt")) + even, odd = T.simd.vdintlv(x0, x1) + lo = T.simd.vmin(even, odd) + hi = T.simd.vmax(even, odd) + lo = T.simd.vmins(lo, valid_limit) + difference = T.simd.vsub(hi, lo, T.simd.vcmps(hi, valid_limit, op="lt")) + T.simd.vsts(smaller[idx0, i * 64], lo) + T.simd.vsts(larger[idx0, i * 64], hi) + if decode: + # Quantized records are contiguous within each page. + # The wrapper currently requires compact page storage. + bytes0 = T.if_then_else(extra0, extra_record_bytes, record_bytes) + byte_offset = T.simd.vmuls(lo, bytes0) + valid_small = T.simd.vcmps(lo, valid_limit, op="lt") + byte_offset = T.simd.vsel(byte_offset, T.simd.vdup(T.uint32(0xFFFFFFFF), "uint32"), valid_small) + T.simd.vsts(offsets[i * 64], byte_offset) + T.simd.vsts(strides[i * 64], T.simd.vmuls(difference, bytes0)) + else: + T.simd.vsts(offsets[i * 64], T.simd.vmuls(lo, 512)) + T.simd.vsts(strides[i * 64], T.simd.vmuls(difference, 1024)) + T.ascend_rls_buf("PIPE_V", 1) + T.ascend_rls_buf("PIPE_V", 0) + # Stage 5: rescale the UB accumulator only when requested. + if step >= 5 and step - 5 < total: + t5 = step - 5 + job5 = T.alloc_var("int32", init=t5 // nb) + k5 = T.alloc_var("int32", init=t5 % nb) + if k5 == 0: + emitted = 0 + rescale5 = T.alloc_var("int32", init=0) + if k5 > 0: + T.ascend_cross_core_wait_flag(4, "PIPE_S", 2) + rescale5 = T.call_extern("int32", "fm_ss_get", (job5 % 8) * 8 + k5 % 8) + if rescale5: + for d in rolled(4): + T.ascend_get_buf("PIPE_V", 4 + d) + T.ascend_cross_core_wait_flag(4, "PIPE_V", 4 + d % 2) + if emitted > 0: + scale_output_vf(True, d, job5, k5, factors, out, accum) + else: + scale_output_vf(False, d, job5, k5, factors, out, accum) + T.ascend_cross_core_set_flag(4, "PIPE_V", 4 + d % 2) + T.ascend_rls_buf("PIPE_V", 4 + d) + emitted = emitted + 1 + if k5 == nb - 1: + for d in T.unroll(4): + T.ascend_get_buf("PIPE_V", 4 + d) + T.ascend_cross_core_wait_flag(4, "PIPE_V", 4 + d % 2) + if emitted > 0: + final_output_vf(True, d, job5, denorm, out, accum, final) + else: + final_output_vf(False, d, job5, denorm, out, accum, final) + T.ascend_cross_core_set_flag(4, "PIPE_V", 4 + d % 2) + T.ascend_rls_buf("PIPE_V", 4 + d) + T.ascend_get_buf("PIPE_MTE3", 4 + d) + T.call_extern( + "void", + "fm_write_o", + T.access_ptr(O[bx + job5 * 32, sid * 32, d * 128], "w"), + T.access_ptr(final[d * 256], "r"), + ) + T.ascend_rls_buf("PIPE_MTE3", 4 + d) + gather_k = step % nb + gather_rel = T.if_then_else(gather_k >= nb0, gather_k - nb0, gather_k) + gather_later = step < total and gather_rel % (ib // 64) == 0 + if gather_later and step >= 4 and step - 4 < total and (step - 4) % nb > 0: + t4 = step - 4 + T.ascend_wait_flag("V_S", 0) + T.call_extern( + "void", "fm_ss_set", (t4 // nb % 8) * 8 + t4 % nb % 8, sid, delta[t4 // nb % 2, t4 % nb % 2, 0] > 6.0 / scale + ) + T.ascend_cross_core_set_flag(4, "PIPE_S", 2) + # Stage 0: 16 two-row DMA instructions per AIV. + if step < total: + k0 = step % nb + extra0 = k0 >= nb0 + rel0 = T.if_then_else(extra0, k0 - nb0, k0) + start = ((rel0 % (ib // 64)) * 64 + sid * 32) // 2 + T.ascend_get_buf("PIPE_MTE2", 8 + step % 3) + T.ascend_get_buf("PIPE_S", 1) + for i in T.unroll(16): + off = offsets[start + i] + if decode: + n_burst = T.if_then_else(off != T.uint32(0xFFFFFFFF), 2, 0) + stride0 = T.Cast("int64", strides[start + i]) + # Clamp the unused source address for fully invalid pairs. + safe_off = T.if_then_else(n_burst == 0, T.uint32(0), off) + # Temporary codegen workaround: omit the copy guards, + # keeping runtime burst counts of 0/2 (see README). + if extra_topk > 0 and extra0: + if extra_fp4: + pair_src = T.make_tensor_from_addr( + T.access_ptr(ExtraKV[safe_off], "r"), (n_burst, 288), "uint8", strides=(stride0, 1) + ) + pair_dst = T.make_tensor_from_addr( + T.access_ptr(raw4[step % 3, i * 2, 0], "w"), + (n_burst, 288), + "uint8", + strides=(288, 1), + storage_scope="shared.dyn", + ) + T.assume(n_burst > 0) + T.copy(pair_src, pair_dst, l2_cache_ctrl="normal_lv" if cache_hint else "notalloc_clean") + else: + pair_src = T.make_tensor_from_addr( + T.access_ptr(ExtraKV[safe_off], "r"), (n_burst, 528), "uint8", strides=(stride0, 1) + ) + pair_dst = T.make_tensor_from_addr( + T.access_ptr(raw8[step % 3, i * 2, 0], "w"), + (n_burst, 528), + "uint8", + strides=(544, 1), + storage_scope="shared.dyn", + ) + T.assume(n_burst > 0) + T.copy(pair_src, pair_dst, l2_cache_ctrl="normal_lv" if cache_hint else "notalloc_clean") + else: + if fp4: + pair_src = T.make_tensor_from_addr( + T.access_ptr(KV[safe_off], "r"), (n_burst, 288), "uint8", strides=(stride0, 1) + ) + pair_dst = T.make_tensor_from_addr( + T.access_ptr(raw4[step % 3, i * 2, 0], "w"), + (n_burst, 288), + "uint8", + strides=(288, 1), + storage_scope="shared.dyn", + ) + T.assume(n_burst > 0) + T.copy(pair_src, pair_dst, l2_cache_ctrl="normal_lv" if cache_hint else "notalloc_clean") + else: + pair_src = T.make_tensor_from_addr( + T.access_ptr(KV[safe_off], "r"), (n_burst, 528), "uint8", strides=(stride0, 1) + ) + pair_dst = T.make_tensor_from_addr( + T.access_ptr(raw8[step % 3, i * 2, 0], "w"), + (n_burst, 528), + "uint8", + strides=(544, 1), + storage_scope="shared.dyn", + ) + T.assume(n_burst > 0) + T.copy(pair_src, pair_dst, l2_cache_ctrl="normal_lv" if cache_hint else "notalloc_clean") + else: + n_burst = T.if_then_else(off != nk * 512, 2, 0) + pair_src = T.make_tensor_from_addr( + T.access_ptr(KV[off // 512, 0], "r"), + (n_burst, 512), + "bfloat16", + strides=(T.Cast("int64", strides[start + i]) // 2, 1), + ) + pair_dst = T.make_tensor_from_addr( + T.access_ptr(kv[step % 3, i * 2, 0], "w"), + (n_burst, 512), + "bfloat16", + strides=(512, 1), + storage_scope="shared.dyn", + ) + # As above, n_burst can still be zero at runtime. + T.assume(n_burst > 0) + T.copy(pair_src, pair_dst, l2_cache_ctrl="normal_lv") + T.ascend_rls_buf("PIPE_S", 1) + T.ascend_rls_buf("PIPE_MTE2", 8 + step % 3) + if not gather_later and step >= 4 and step - 4 < total and (step - 4) % nb > 0: + t4 = step - 4 + T.ascend_wait_flag("V_S", 0) + T.call_extern( + "void", "fm_ss_set", (t4 // nb % 8) * 8 + t4 % nb % 8, sid, delta[t4 // nb % 2, t4 % nb % 2, 0] > 6.0 / scale + ) + T.ascend_cross_core_set_flag(4, "PIPE_S", 2) + T.ascend_pipe_barrier("PIPE_ALL") + + return main diff --git a/examples/ascend/flashmla/decode.py b/examples/ascend/flashmla/decode.py new file mode 100644 index 0000000000..dad0be9634 --- /dev/null +++ b/examples/ascend/flashmla/decode.py @@ -0,0 +1,103 @@ +"""Explicitly scheduled FP8 main cache / FP8 or FP4 extra cache attention.""" + +import torch +import tilelang +from tilelang import tvm + + +def prepare_decode(param, inputs, *, implementation="manual"): + """Prepare the PR's packed KV inputs without changing their quantization.""" + if implementation == "manual": + from examples.ascend.flashmla.core import PASS_CONFIG, sparse_attention + elif implementation == "staged": + from examples.ascend.flashmla.staged import PASS_CONFIG, sparse_attention + else: + raise ValueError(f"Unknown implementation: {implementation}") + nq = param.decode.b * param.s_q + q = inputs.q.reshape(nq, 64, 512) + main = inputs.kv_scope + extra = inputs.extra_kv_scope + kv = main.get_kvcache_for_flash_mla() + if not kv.is_contiguous(): + raise ValueError("The example currently requires contiguous packed KV pages") + indices = main.indices_in_kvcache.reshape(nq, param.topk) + sink = inputs.attn_sink if inputs.attn_sink is not None else torch.zeros(64, device=q.device, dtype=torch.float32) + lengths = ( + main.topk_length.repeat_interleave(param.s_q) + if main.topk_length is not None + else torch.full((nq,), param.topk, device=q.device, dtype=torch.int32) + ) + if extra is not None: + ekv = extra.get_kvcache_for_flash_mla() + if not ekv.is_contiguous(): + raise ValueError("The example currently requires contiguous packed extra KV pages") + etopk = extra.indices_in_kvcache.shape[-1] + eindices = extra.indices_in_kvcache.reshape(nq, etopk) + elengths = ( + extra.topk_length.repeat_interleave(param.s_q) + if extra.topk_length is not None + else torch.full((nq,), etopk, device=q.device, dtype=torch.int32) + ) + eformat = "fp4" if ekv.shape[-1] == 288 else "fp8" + else: + ekv = torch.empty(1, device=q.device, dtype=torch.uint8) + etopk = 0 + eindices = torch.zeros((nq, 1), device=q.device, dtype=torch.int32) + elengths = torch.zeros(nq, device=q.device, dtype=torch.int32) + eformat = "fp8" + nbytes = kv.numel() + (ekv.numel() if extra is not None else 0) + ntokens = kv.numel() // 528 + (ekv.numel() // (288 if eformat == "fp4" else 528) if extra is not None else 0) + reuse_cache = ( + nbytes <= 8 * 1024**2 + and nq * (param.topk + etopk) >= 4 * ntokens + and main.topk_length is None + and (extra is None or extra.topk_length is None) + ) + func = sparse_attention( + nq, + ntokens, + param.topk, + sink=param.have_attn_sink, + variable_lengths=main.topk_length is not None, + scale=inputs.sm_scale, + q_strides=tuple(q.stride()), + index_stride=indices.stride(0), + kv_format="fp8", + extra_topk=etopk, + extra_format=eformat, + extra_index_stride=eindices.stride(0), + variable_extra_lengths=extra is not None and extra.topk_length is not None, + kv_storage_bytes=kv.numel(), + extra_storage_bytes=ekv.numel(), + cache_hint=int(reuse_cache), + ) + if implementation == "manual": + params = list(func.params[:9]) + list(func.params[10:]) + param_set = set(params) + func = tvm.tirx.PrimFunc( + params, func.body, func.ret_type, {var: buf for var, buf in func.buffer_map.items() if var in param_set}, func.attrs + ) + kernel = tilelang.compile( + func, + target="ascend", + out_idx=[8, 9], + pass_configs=PASS_CONFIG, + compile_flags=["-Ofast", "-mllvm", "-enable-hiipu-vf-loop-unroll"], + ) + kv_bytes = kv.view(torch.uint8).view(-1) + extra_bytes = ekv.view(torch.uint8).view(-1) + + def run(): + # Match the PR wrapper's per-invocation flattening. Non-contiguous + # [batch, query, ...] tensors can require a device copy here, which + # affects the attention kernel's initial cache state even when the + # profiler only times the attention kernel. + query = inputs.q.reshape(nq, 64, 512) + main_indices = main.indices_in_kvcache.reshape(nq, param.topk) + extra_indices = extra.indices_in_kvcache.reshape(nq, etopk) if extra is not None else eindices + main_lengths = main.topk_length.repeat_interleave(param.s_q) if main.topk_length is not None else lengths + extra_lengths = extra.topk_length.repeat_interleave(param.s_q) if extra is not None and extra.topk_length is not None else elengths + out, lse = kernel(query, kv_bytes, main_indices, sink, main_lengths, extra_bytes, extra_indices, extra_lengths) + return out.reshape(param.decode.b, param.s_q, 64, 512), lse.reshape(param.decode.b, param.s_q, 64).transpose(1, 2).contiguous() + + return kernel, run diff --git a/examples/ascend/flashmla/manual.py b/examples/ascend/flashmla/manual.py new file mode 100644 index 0000000000..00fbcaa041 --- /dev/null +++ b/examples/ascend/flashmla/manual.py @@ -0,0 +1,26 @@ +"""BF16 prefill entry point for the explicitly scheduled sparse attention core.""" + +import tilelang +from tilelang import tvm +from examples.ascend.flashmla.core import PASS_CONFIG, sparse_attention + + +def prefill(nq, nk, topk, **kwargs): + func = sparse_attention(nq, nk, topk, **kwargs) + # The BF16 specialization never references the three decode-only inputs. + # Remove them from its public ABI; the pipeline implementation stays shared. + params = list(func.params[:5]) + list(func.params[8:]) + param_set = set(params) + return tvm.tirx.PrimFunc( + params, func.body, func.ret_type, {var: buf for var, buf in func.buffer_map.items() if var in param_set}, func.attrs + ) + + +def compile_prefill(nq, nk, topk, **kwargs): + return tilelang.compile( + prefill(nq, nk, topk, **kwargs), + target="ascend", + out_idx=[5, 6, 7], + pass_configs=PASS_CONFIG, + compile_flags=["-Ofast", "-mllvm", "-enable-hiipu-vf-loop-unroll"], + ) diff --git a/examples/ascend/flashmla/manual_intrinsics.h b/examples/ascend/flashmla/manual_intrinsics.h new file mode 100644 index 0000000000..d288b16bdb --- /dev/null +++ b/examples/ascend/flashmla/manual_intrinsics.h @@ -0,0 +1,64 @@ +// Low-level instruction adapters for the manually scheduled FlashMLA example. +// Algorithm and memory plan: deepseek-ai/FlashMLA PR #229 (MIT license). +#pragma once +#include "c_api/asc_simd.h" +#include "kernel_operator.h" + +__aicore__ inline void fm_init() { + AscendC::InitSocState(); + asc_set_gm2ub_loop_size(1ul, 1ul); + asc_set_ub2gm_loop_size(1ul, 1ul); + asc_set_ctrl(0ull); + asc_set_gm2l1_nz_para(uint64_t(1) | (uint64_t(1) << 16) | + (uint64_t(32) << 32)); + asc_set_l0c2gm_nz2nd(1u, 0u, 0u); + asc_set_l3d_rpt_b(1 << 16); +} +__aicore__ inline void fm_ss_set(uint32_t slot, uint32_t sid, uint32_t value) { + auto p = (volatile __ssbuf__ uint32_t *)AscendC::GetSsbufBaseAddr(); + p[slot * 2 + sid] = value; +} +__aicore__ inline uint32_t fm_ss_get(uint32_t slot) { + auto p = (volatile __ssbuf__ uint64_t *)AscendC::GetSsbufBaseAddr(); + return p[slot] != 0; +} +__aicore__ inline void fm_q_load(__cbuf__ bfloat16_t *dst, + __gm__ bfloat16_t *src, uint32_t stride) { + asc_copy_gm2l1_nd2nz(dst, src, stride, 0u, 32u, 512u, 0u, false); +} +__aicore__ inline void fm_kv_push(__cbuf__ bfloat16_t *dst, + __ubuf__ bfloat16_t *src) { + asc_copy_ub2l1(dst, src, 32u, 32u, 1u, 0u); +} +__aicore__ inline void fm_a_load(__ca__ bfloat16_t *dst, + __cbuf__ bfloat16_t *src, uint32_t k) { + asc_copy_l12l0a(dst, src, 0, k * 8u, 2u, 8u, 2u, 4u); +} +__aicore__ inline void fm_k_load(__cb__ bfloat16_t *dst, + __cbuf__ bfloat16_t *src, uint32_t k) { + asc_copy_l12l0b(dst, src, 0, k * 8u, 2u, 8u, 2u, 4u); +} +__aicore__ inline void fm_v_load(__cb__ bfloat16_t *dst, + __cbuf__ bfloat16_t *src, uint32_t k) { + asc_copy_l12l0b_transpose(dst, src, 0, k * 8u, 2u, 8u, 2u, 8u); +} +__aicore__ inline void fm_s_load(__ca__ bfloat16_t *dst, + __cbuf__ bfloat16_t *src) { + asc_copy_l12l0a(dst, src, 0, 0, 4u, 4u, 4u, 4u); +} +__aicore__ inline void fm_mmad(__cc__ float *dst, __ca__ bfloat16_t *a, + __cb__ bfloat16_t *b, uint32_t n, uint32_t k, + uint32_t unit, bool clear) { + asc_mmad(dst, a, b, 64u, k, n, unit, true, false, clear); +} +__aicore__ inline void fm_o_store(__ubuf__ float *dst, __cc__ float *src, + uint32_t unit) { + asc_copy_l0c2ub(dst, src, 128u, 64u, 128u, 64u, 1, false, 0, unit, 0, 0, + false, true, 0, 0, 0, 0, 0, 0, 0, 0); +} +__aicore__ inline void fm_write_o(__gm__ bfloat16_t *dst, + __ubuf__ bfloat16_t *src) { + asc_copy_ub2gm_align(dst, src, 32u, 256u, + asc_store_l2_cache_mode::NORMAL_FIRST_VICTIM, 1024ul, + 2048ul); +} diff --git a/examples/ascend/flashmla/quantization.py b/examples/ascend/flashmla/quantization.py new file mode 100644 index 0000000000..490935fb97 --- /dev/null +++ b/examples/ascend/flashmla/quantization.py @@ -0,0 +1,64 @@ +"""SIMD dequantization of the interleaved V4.1 KV records from PR #229.""" + +import tilelang.ascend.language as T + + +def rolled(n): + return T.unroll(n, annotations={"pragma_unroll_factor": 1}) + + +@T.macro +def dequantize(raw, scales, nz, slot, fp4): + with T.SimdVF(): + m16 = T.simd.pset(16) + m8 = T.simd.pset(8) + arange = T.reinterpret("uint16x128", T.simd.vci(0, "int16")) + if fp4: + row_indices = T.simd.vshls(T.simd.vshrs(arange, 5), 8) + scale_indices = T.simd.vadd(arange, row_indices) + for row in rolled(8): + packed_bits = T.simd.vgather2((raw[slot, row * 4, 256] if slot is not None else raw[row * 4, 256]), scale_indices, m16) + packed = T.reinterpret("float8_e4m3fnx256", packed_bits) + even = T.simd.vcvt(packed, "float32", m16, part=0) + odd = T.simd.vcvt(packed, "float32", m16, part=2) + even_bits = T.simd.vshrs(T.reinterpret("uint32x64", even), 16) + bits = T.simd.vor(even_bits, T.reinterpret("uint32x64", odd)) + T.simd.vsts(scales[row * 128], T.reinterpret("bfloat16x128", bits), m16) + else: + row_indices = T.simd.vmuls(T.simd.vshrs(arange, 4), 528) + scale_indices = T.simd.vadd(arange, row_indices) + for row in rolled(4): + packed = T.simd.vgather2((raw[slot, row * 8, 512] if slot is not None else raw[row * 8, 512]), scale_indices, m16) + clamped = T.simd.vmins(packed, 0xE0) + bits = T.simd.vshls(clamped, 7) + T.simd.vsts(scales[row * 128], T.reinterpret("bfloat16x128", bits), m16) + T.simd.mem_bar("VST_VLD") + for row in rolled(32): + sf = T.alloc_local((4,), "bfloat16x128") + if not fp4: + sf01 = T.simd.vld(scales[row * 16], "E2B_B16") + sf23 = T.simd.vld(scales[row * 16 + 8], "E2B_B16") + s0, s1 = T.simd.vintlv(sf01, sf01) + s2, s3 = T.simd.vintlv(sf23, sf23) + sf[0] = s0 + sf[1] = s1 + sf[2] = s2 + sf[3] = s3 + for chunk in T.unroll(4): + if fp4: + scale = T.simd.vld(scales[row * 32 + chunk * 8], "E2B_B16") + bits = T.simd.vld((raw[slot, row, chunk * 64] if slot is not None else raw[row, chunk * 64]), "UNPK4_B8") + packed = T.reinterpret("float4_e2m1fnx512", bits) + value = T.simd.vcvt(packed, "bfloat16", m16, part=0) + dequant = T.simd.vmul(value, scale, m16) + T.simd.vsstb(dequant, (nz[chunk * 4224 + row * 16] if slot is not None else nz[row, chunk * 128]), 33 << 16, m16) + else: + bits = T.simd.vld((raw[slot, row, chunk * 128] if slot is not None else raw[row, chunk * 128]), "UNPK_B8") + packed = T.reinterpret("float8_e4m3fnx256", bits) + even = T.simd.vcvt(packed, "float32", m8, part=0) + odd = T.simd.vcvt(packed, "float32", m8, part=2) + even_bits = T.simd.vshrs(T.reinterpret("uint32x64", even), 16) + data_bits = T.simd.vor(even_bits, T.reinterpret("uint32x64", odd)) + value = T.reinterpret("bfloat16x128", data_bits) + dequant = T.simd.vmul(value, sf[chunk], m16) + T.simd.vsstb(dequant, (nz[chunk * 4224 + row * 16] if slot is not None else nz[row, chunk * 128]), 33 << 16, m16) diff --git a/examples/ascend/flashmla/results/benchmark-input-metadata.json b/examples/ascend/flashmla/results/benchmark-input-metadata.json new file mode 100644 index 0000000000..029de8a363 --- /dev/null +++ b/examples/ascend/flashmla/results/benchmark-input-metadata.json @@ -0,0 +1,675 @@ +{ + "provenance": "Shapes reconstructed from pinned lib.py/quant.py/generate.py with seed=42; flattened strides verified in each saved ASC.", + "cases": [ + { + "case": 0, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 64, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 256, + "max_padded_sequence": 1024, + "shape": [ + 256, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 34603008 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 4096, + "max_padded_sequence": 4096, + "shape": [ + 4096, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 138412032 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 1, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 128, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 512, + "max_padded_sequence": 1024, + "shape": [ + 512, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 69206016 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 11264, + "max_padded_sequence": 5632, + "shape": [ + 11264, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 380633088 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 2, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 256, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 1024, + "max_padded_sequence": 1024, + "shape": [ + 1024, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 138412032 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 22528, + "max_padded_sequence": 5632, + "shape": [ + 22528, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 761266176 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 3, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 512, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 2048, + "max_padded_sequence": 1024, + "shape": [ + 2048, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 276824064 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 45056, + "max_padded_sequence": 5632, + "shape": [ + 45056, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 1522532352 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 4, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 64, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 256, + "max_padded_sequence": 1024, + "shape": [ + 256, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 34603008 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 4096, + "max_padded_sequence": 4096, + "shape": [ + 4096, + 64, + 1, + 288 + ], + "strides": [ + 18432, + 288, + 288, + 1 + ], + "storage_bytes": 75497472 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 5, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 128, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 512, + "max_padded_sequence": 1024, + "shape": [ + 512, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 69206016 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 11264, + "max_padded_sequence": 5632, + "shape": [ + 11264, + 64, + 1, + 288 + ], + "strides": [ + 18432, + 288, + 288, + 1 + ], + "storage_bytes": 207618048 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 6, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 256, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 1024, + "max_padded_sequence": 1024, + "shape": [ + 1024, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 138412032 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 22528, + "max_padded_sequence": 5632, + "shape": [ + 22528, + 64, + 1, + 288 + ], + "strides": [ + 18432, + 288, + 288, + 1 + ], + "storage_bytes": 415236096 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 7, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 512, + 4, + 64, + 512 + ], + "query_original_strides": [ + 208000, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 32768, + 512, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 128, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 512, + 1 + ], + "main_cache": { + "block_size": 256, + "num_blocks": 2048, + "max_padded_sequence": 1024, + "shape": [ + 2048, + 256, + 1, + 528 + ], + "strides": [ + 135168, + 528, + 528, + 1 + ], + "storage_bytes": 276824064 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 45056, + "max_padded_sequence": 5632, + "shape": [ + 45056, + 64, + 1, + 288 + ], + "strides": [ + 18432, + 288, + 288, + 1 + ], + "storage_bytes": 830472192 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NOTALLOC_CLEAN" + }, + { + "case": 8, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 1, + 4096, + 64, + 512 + ], + "query_original_strides": [ + 170435200, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 41600, + 640, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 256, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 640, + 1 + ], + "main_cache": { + "block_size": 64, + "num_blocks": 64, + "max_padded_sequence": 4096, + "shape": [ + 64, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 2162688 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 64, + "max_padded_sequence": 4096, + "shape": [ + 64, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 2162688 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NORMAL_LAST_VICTIM" + }, + { + "case": 9, + "sm_scale": 0.032352028870043036, + "query_shape": [ + 1, + 4096, + 64, + 512 + ], + "query_original_strides": [ + 170435200, + 41600, + 640, + 1 + ], + "flattened_query_strides_verified_in_asc": [ + 41600, + 640, + 1 + ], + "flattened_main_index_strides_verified_in_asc": [ + 256, + 1 + ], + "flattened_extra_index_strides_verified_in_asc": [ + 640, + 1 + ], + "main_cache": { + "block_size": 64, + "num_blocks": 64, + "max_padded_sequence": 4096, + "shape": [ + 64, + 64, + 1, + 528 + ], + "strides": [ + 33792, + 528, + 528, + 1 + ], + "storage_bytes": 2162688 + }, + "extra_cache": { + "block_size": 64, + "num_blocks": 64, + "max_padded_sequence": 4096, + "shape": [ + 64, + 64, + 1, + 288 + ], + "strides": [ + 18432, + 288, + 288, + 1 + ], + "storage_bytes": 1179648 + }, + "cache_container_dtype": "float8_e4m3fn (manual kernel views raw bytes as uint8)", + "cache_hint": "NORMAL_LAST_VICTIM" + } + ] +} diff --git a/examples/ascend/flashmla/results/decode-bench-final.json b/examples/ascend/flashmla/results/decode-bench-final.json new file mode 100644 index 0000000000..6850579463 --- /dev/null +++ b/examples/ascend/flashmla/results/decode-bench-final.json @@ -0,0 +1,1088 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "id": 0, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 78.43971252441406, + 78.3681869506836, + 78.98807525634766 + ], + "launch_us": [ + [ + 78.2012939453125, + 76.2939453125, + 78.43971252441406, + 78.43971252441406, + 79.39338684082031, + 79.63180541992188, + 78.2012939453125, + 78.2012939453125, + 78.91654968261719, + 78.67813110351562 + ], + [ + 79.87022399902344, + 79.63180541992188, + 77.48603820800781, + 78.91654968261719, + 78.67813110351562, + 78.2012939453125, + 77.24761962890625, + 78.67813110351562, + 77.72445678710938, + 77.24761962890625 + ], + [ + 80.108642578125, + 79.39338684082031, + 77.48603820800781, + 78.2012939453125, + 78.2012939453125, + 80.82389831542969, + 79.39338684082031, + 78.67813110351562, + 78.67813110351562, + 78.91654968261719 + ] + ], + "median_us": 78.43971252441406, + "tflops": 273.7750533355925 + }, + "manual": { + "round_mean_us": [ + 76.65157318115234, + 76.69925689697266, + 76.96151733398438 + ], + "launch_us": [ + [ + 78.2012939453125, + 77.00920104980469, + 77.00920104980469, + 76.05552673339844, + 75.57868957519531, + 75.81710815429688, + 76.53236389160156, + 77.96287536621094, + 76.53236389160156, + 75.81710815429688 + ], + [ + 77.00920104980469, + 77.24761962890625, + 76.2939453125, + 77.00920104980469, + 76.53236389160156, + 76.53236389160156, + 77.72445678710938, + 76.05552673339844, + 76.2939453125, + 76.2939453125 + ], + [ + 76.05552673339844, + 77.00920104980469, + 77.72445678710938, + 76.2939453125, + 76.2939453125, + 77.96287536621094, + 77.00920104980469, + 77.00920104980469, + 77.24761962890625, + 77.00920104980469 + ] + ], + "median_us": 76.69925689697266, + "tflops": 279.98754288905786 + } + } + }, + { + "id": 1, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 144.29092407226562, + 144.45781707763672, + 144.9108123779297 + ], + "launch_us": [ + [ + 143.28956604003906, + 144.95849609375, + 144.72007751464844, + 144.72007751464844, + 143.7664031982422, + 144.95849609375, + 144.72007751464844, + 144.2432403564453, + 142.81272888183594, + 144.72007751464844 + ], + [ + 144.72007751464844, + 145.19691467285156, + 144.72007751464844, + 143.0511474609375, + 144.95849609375, + 144.95849609375, + 145.43533325195312, + 143.7664031982422, + 144.2432403564453, + 143.52798461914062 + ], + [ + 145.6737518310547, + 144.72007751464844, + 144.48165893554688, + 144.95849609375, + 144.00482177734375, + 144.48165893554688, + 145.91217041015625, + 144.95849609375, + 145.19691467285156, + 144.72007751464844 + ] + ], + "median_us": 144.45781707763672, + "tflops": 297.31636424297716 + }, + "manual": { + "round_mean_us": [ + 141.9544219970703, + 142.38357543945312, + 142.3358917236328 + ], + "launch_us": [ + [ + 140.90538024902344, + 142.57431030273438, + 142.57431030273438, + 141.62063598632812, + 142.09747314453125, + 142.3358917236328, + 143.7664031982422, + 140.90538024902344, + 141.143798828125, + 141.62063598632812 + ], + [ + 142.57431030273438, + 142.09747314453125, + 142.3358917236328, + 140.90538024902344, + 142.81272888183594, + 142.3358917236328, + 142.09747314453125, + 143.0511474609375, + 143.0511474609375, + 142.57431030273438 + ], + [ + 142.09747314453125, + 143.28956604003906, + 142.57431030273438, + 141.143798828125, + 143.52798461914062, + 142.57431030273438, + 142.3358917236328, + 142.81272888183594, + 141.8590545654297, + 141.143798828125 + ] + ], + "median_us": 142.3358917236328, + "tflops": 301.748718751792 + } + } + }, + { + "id": 2, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 274.8250961303711, + 274.51515197753906, + 274.96814727783203 + ], + "launch_us": [ + [ + 275.3734588623047, + 274.89662170410156, + 275.8502960205078, + 274.89662170410156, + 274.41978454589844, + 274.89662170410156, + 274.41978454589844, + 274.41978454589844, + 274.1813659667969, + 274.89662170410156 + ], + [ + 274.41978454589844, + 274.1813659667969, + 274.1813659667969, + 274.89662170410156, + 274.41978454589844, + 273.2276916503906, + 275.8502960205078, + 274.89662170410156, + 274.1813659667969, + 274.89662170410156 + ], + [ + 275.61187744140625, + 274.658203125, + 275.3734588623047, + 275.8502960205078, + 274.1813659667969, + 274.41978454589844, + 274.41978454589844, + 273.70452880859375, + 275.1350402832031, + 276.32713317871094 + ] + ], + "median_us": 274.8250961303711, + "tflops": 312.5600504811657 + }, + "manual": { + "round_mean_us": [ + 271.29650115966797, + 271.24881744384766, + 271.3441848754883 + ], + "launch_us": [ + [ + 270.3666687011719, + 270.3666687011719, + 271.79718017578125, + 272.2740173339844, + 269.88983154296875, + 270.843505859375, + 271.79718017578125, + 271.3203430175781, + 272.2740173339844, + 272.0355987548828 + ], + [ + 270.843505859375, + 271.08192443847656, + 271.3203430175781, + 271.79718017578125, + 271.3203430175781, + 272.0355987548828, + 271.3203430175781, + 270.3666687011719, + 271.79718017578125, + 270.60508728027344 + ], + [ + 272.51243591308594, + 271.79718017578125, + 270.1282501220703, + 270.843505859375, + 270.60508728027344, + 271.5587615966797, + 271.79718017578125, + 271.08192443847656, + 271.5587615966797, + 271.5587615966797 + ] + ], + "median_us": 271.29650115966797, + "tflops": 316.62533631218884 + } + } + }, + { + "id": 3, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 540.0419235229492, + 539.1836166381836, + 539.4458770751953 + ], + "launch_us": [ + [ + 541.4485931396484, + 539.3028259277344, + 538.1107330322266, + 540.4949188232422, + 539.7796630859375, + 539.5412445068359, + 540.4949188232422, + 540.4949188232422, + 540.4949188232422, + 540.2565002441406 + ], + [ + 538.1107330322266, + 541.2101745605469, + 539.3028259277344, + 540.0180816650391, + 538.5875701904297, + 540.2565002441406, + 538.8259887695312, + 538.5875701904297, + 538.5875701904297, + 538.3491516113281 + ], + [ + 539.7796630859375, + 539.7796630859375, + 539.5412445068359, + 539.0644073486328, + 539.5412445068359, + 540.2565002441406, + 540.4949188232422, + 538.3491516113281, + 537.872314453125, + 539.7796630859375 + ] + ], + "median_us": 539.4458770751953, + "tflops": 318.47252734874894 + }, + "manual": { + "round_mean_us": [ + 534.0814590454102, + 533.4138870239258, + 533.5569381713867 + ], + "launch_us": [ + [ + 534.7728729248047, + 534.2960357666016, + 533.5807800292969, + 532.6271057128906, + 534.0576171875, + 534.5344543457031, + 535.2497100830078, + 535.0112915039062, + 532.8655242919922, + 533.8191986083984 + ], + [ + 535.7265472412109, + 533.1039428710938, + 533.3423614501953, + 533.3423614501953, + 532.6271057128906, + 533.1039428710938, + 533.5807800292969, + 534.0576171875, + 532.3886871337891, + 532.8655242919922 + ], + [ + 535.7265472412109, + 532.8655242919922, + 532.6271057128906, + 533.3423614501953, + 533.3423614501953, + 534.5344543457031, + 534.2960357666016, + 533.3423614501953, + 532.3886871337891, + 533.1039428710938 + ] + ], + "median_us": 533.5569381713867, + "tflops": 321.9875509983822 + } + } + }, + { + "id": 4, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 70.66726684570312, + 70.59574127197266, + 70.50037384033203 + ], + "launch_us": [ + [ + 69.85664367675781, + 72.00241088867188, + 70.81031799316406, + 71.28715515136719, + 71.76399230957031, + 69.37980651855469, + 71.76399230957031, + 70.5718994140625, + 69.37980651855469, + 69.85664367675781 + ], + [ + 69.85664367675781, + 71.04873657226562, + 72.00241088867188, + 69.14138793945312, + 68.90296936035156, + 69.37980651855469, + 69.37980651855469, + 71.28715515136719, + 72.71766662597656, + 72.24082946777344 + ], + [ + 70.81031799316406, + 70.09506225585938, + 69.37980651855469, + 70.09506225585938, + 69.37980651855469, + 72.24082946777344, + 69.85664367675781, + 71.28715515136719, + 70.09506225585938, + 71.76399230957031 + ] + ], + "median_us": 70.59574127197266, + "tflops": 304.1945037062138 + }, + "manual": { + "round_mean_us": [ + 71.38252258300781, + 70.88184356689453, + 71.5494155883789 + ], + "launch_us": [ + [ + 70.09506225585938, + 72.95608520507812, + 72.95608520507812, + 70.09506225585938, + 71.76399230957031, + 71.04873657226562, + 70.33348083496094, + 72.24082946777344, + 70.09506225585938, + 72.24082946777344 + ], + [ + 70.33348083496094, + 71.52557373046875, + 72.71766662597656, + 70.33348083496094, + 71.76399230957031, + 70.81031799316406, + 70.5718994140625, + 70.33348083496094, + 70.5718994140625, + 69.85664367675781 + ], + [ + 70.33348083496094, + 72.95608520507812, + 71.52557373046875, + 70.81031799316406, + 71.76399230957031, + 70.5718994140625, + 73.43292236328125, + 70.81031799316406, + 70.09506225585938, + 73.19450378417969 + ] + ], + "median_us": 71.38252258300781, + "tflops": 300.84165847498303 + } + } + }, + { + "id": 5, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 131.34479522705078, + 132.3223114013672, + 131.34479522705078 + ], + "launch_us": [ + [ + 130.8917999267578, + 132.08389282226562, + 129.93812561035156, + 132.3223114013672, + 131.13021850585938, + 132.08389282226562, + 131.6070556640625, + 129.22286987304688, + 131.13021850585938, + 133.03756713867188 + ], + [ + 132.3223114013672, + 131.13021850585938, + 131.36863708496094, + 133.514404296875, + 133.27598571777344, + 131.6070556640625, + 133.03756713867188, + 132.7991485595703, + 131.6070556640625, + 132.56072998046875 + ], + [ + 133.03756713867188, + 129.46128845214844, + 130.8917999267578, + 132.08389282226562, + 131.6070556640625, + 130.17654418945312, + 130.17654418945312, + 131.36863708496094, + 131.84547424316406, + 132.7991485595703 + ] + ], + "median_us": 131.34479522705078, + "tflops": 326.9994283805043 + }, + "manual": { + "round_mean_us": [ + 131.988525390625, + 171.75674438476562, + 131.79779052734375 + ], + "launch_us": [ + [ + 131.6070556640625, + 131.36863708496094, + 132.7991485595703, + 131.13021850585938, + 132.08389282226562, + 133.514404296875, + 132.08389282226562, + 131.13021850585938, + 133.03756713867188, + 131.13021850585938 + ], + [ + 152.11105346679688, + 176.1913299560547, + 173.5687255859375, + 173.5687255859375, + 176.42974853515625, + 176.90658569335938, + 178.81393432617188, + 164.031982421875, + 172.13821411132812, + 173.80714416503906 + ], + [ + 132.3223114013672, + 132.08389282226562, + 130.4149627685547, + 133.514404296875, + 132.56072998046875, + 131.36863708496094, + 130.65338134765625, + 133.03756713867188, + 130.8917999267578, + 131.13021850585938 + ] + ], + "median_us": 131.988525390625, + "tflops": 325.4045973533595 + } + } + }, + { + "id": 6, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 246.9062805175781, + 247.43080139160156, + 247.47848510742188 + ], + "launch_us": [ + [ + 247.00164794921875, + 246.7632293701172, + 246.0479736328125, + 247.2400665283203, + 248.19374084472656, + 246.28639221191406, + 245.3327178955078, + 248.43215942382812, + 246.52481079101562, + 247.2400665283203 + ], + [ + 247.2400665283203, + 247.2400665283203, + 246.0479736328125, + 247.47848510742188, + 248.6705780029297, + 247.71690368652344, + 249.38583374023438, + 247.00164794921875, + 247.2400665283203, + 246.28639221191406 + ], + [ + 246.52481079101562, + 248.19374084472656, + 247.955322265625, + 247.955322265625, + 248.90899658203125, + 247.71690368652344, + 247.2400665283203, + 245.80955505371094, + 248.19374084472656, + 246.28639221191406 + ] + ], + "median_us": 247.43080139160156, + "tflops": 347.1651283384464 + }, + "manual": { + "round_mean_us": [ + 248.45600128173825, + 248.05068969726565, + 248.0983734130859 + ], + "launch_us": [ + [ + 247.955322265625, + 248.90899658203125, + 248.19374084472656, + 249.1474151611328, + 247.955322265625, + 248.43215942382812, + 249.1474151611328, + 249.1474151611328, + 247.2400665283203, + 248.43215942382812 + ], + [ + 247.47848510742188, + 248.43215942382812, + 247.955322265625, + 249.62425231933594, + 248.43215942382812, + 247.71690368652344, + 247.71690368652344, + 248.19374084472656, + 247.71690368652344, + 247.2400665283203 + ], + [ + 247.47848510742188, + 247.00164794921875, + 248.6705780029297, + 247.2400665283203, + 249.1474151611328, + 247.47848510742188, + 247.955322265625, + 249.8626708984375, + 246.52481079101562, + 249.62425231933594 + ] + ], + "median_us": 248.0983734130859, + "tflops": 346.23099191777794 + } + } + }, + { + "id": 7, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 490.0932312011719, + 488.6627197265625, + 489.6640777587891 + ], + "launch_us": [ + [ + 490.1885986328125, + 488.9965057373047, + 491.14227294921875, + 491.3806915283203, + 488.28125, + 488.51966857910156, + 487.0891571044922, + 491.14227294921875, + 494.0032958984375, + 490.1885986328125 + ], + [ + 491.6191101074219, + 489.4733428955078, + 490.1885986328125, + 486.3739013671875, + 489.95018005371094, + 488.28125, + 486.3739013671875, + 488.28125, + 488.7580871582031, + 487.32757568359375 + ], + [ + 491.6191101074219, + 488.9965057373047, + 489.4733428955078, + 490.6654357910156, + 488.7580871582031, + 487.8044128417969, + 491.6191101074219, + 490.1885986328125, + 490.42701721191406, + 487.0891571044922 + ] + ], + "median_us": 489.6640777587891, + "tflops": 350.8501024341607 + }, + "manual": { + "round_mean_us": [ + 497.6272583007812, + 497.19810485839844, + 496.3874816894531 + ], + "launch_us": [ + [ + 495.4338073730469, + 499.96376037597656, + 494.7185516357422, + 499.96376037597656, + 499.0100860595703, + 498.77166748046875, + 495.67222595214844, + 498.77166748046875, + 497.8179931640625, + 496.14906311035156 + ], + [ + 498.77166748046875, + 495.67222595214844, + 501.1558532714844, + 496.6259002685547, + 496.14906311035156, + 497.1027374267578, + 496.86431884765625, + 497.57957458496094, + 497.8179931640625, + 494.24171447753906 + ], + [ + 495.67222595214844, + 493.04962158203125, + 494.95697021484375, + 499.96376037597656, + 496.6259002685547, + 498.77166748046875, + 495.4338073730469, + 497.3411560058594, + 494.4801330566406, + 497.57957458496094 + ] + ], + "median_us": 497.19810485839844, + "tflops": 345.5336819695403 + } + } + }, + { + "id": 8, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 1045.536994934082, + 1046.609878540039, + 1045.8946228027344 + ], + "launch_us": [ + [ + 1044.7502136230469, + 1045.7038879394531, + 1044.9886322021484, + 1045.4654693603516, + 1044.9886322021484, + 1044.9886322021484, + 1045.7038879394531, + 1046.895980834961, + 1045.9423065185547, + 1045.9423065185547 + ], + [ + 1044.7502136230469, + 1047.372817993164, + 1047.1343994140625, + 1046.895980834961, + 1046.4191436767578, + 1046.6575622558594, + 1045.9423065185547, + 1047.372817993164, + 1046.4191436767578, + 1047.1343994140625 + ], + [ + 1046.1807250976562, + 1047.372817993164, + 1045.9423065185547, + 1046.1807250976562, + 1047.1343994140625, + 1044.2733764648438, + 1044.9886322021484, + 1046.1807250976562, + 1044.5117950439453, + 1046.1807250976562 + ] + ], + "median_us": 1045.8946228027344, + "tflops": 328.52007858998786 + }, + "manual": { + "round_mean_us": [ + 1032.7577590942383, + 1032.6862335205078, + 1032.5431823730469 + ], + "launch_us": [ + [ + 1033.5445404052734, + 1032.5908660888672, + 1031.637191772461, + 1034.2597961425781, + 1031.637191772461, + 1030.9219360351562, + 1033.3061218261719, + 1034.9750518798828, + 1032.8292846679688, + 1031.8756103515625 + ], + [ + 1034.7366333007812, + 1032.114028930664, + 1032.114028930664, + 1031.1603546142578, + 1031.8756103515625, + 1031.3987731933594, + 1034.2597961425781, + 1033.782958984375, + 1033.0677032470703, + 1032.3524475097656 + ], + [ + 1034.7366333007812, + 1032.5908660888672, + 1032.114028930664, + 1032.3524475097656, + 1031.8756103515625, + 1033.782958984375, + 1031.1603546142578, + 1033.3061218261719, + 1032.114028930664, + 1031.3987731933594 + ] + ], + "median_us": 1032.6862335205078, + "tflops": 332.7219561247076 + } + } + }, + { + "id": 9, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 936.8419647216797, + 935.7452392578125, + 939.0592575073242 + ], + "launch_us": [ + [ + 933.8855743408203, + 940.3228759765625, + 931.7398071289062, + 945.3296661376953, + 934.600830078125, + 940.5612945556641, + 934.600830078125, + 935.0776672363281, + 938.41552734375, + 933.8855743408203 + ], + [ + 934.8392486572266, + 943.8991546630859, + 934.8392486572266, + 934.1239929199219, + 932.2166442871094, + 937.4618530273438, + 935.0776672363281, + 935.7929229736328, + 934.600830078125, + 934.600830078125 + ], + [ + 934.3624114990234, + 932.2166442871094, + 935.7929229736328, + 949.3827819824219, + 938.1771087646484, + 932.4550628662109, + 949.6212005615234, + 946.7601776123047, + 935.3160858154297, + 936.5081787109375 + ] + ], + "median_us": 936.8419647216797, + "tflops": 366.7613072628286 + }, + "manual": { + "round_mean_us": [ + 6749.844551086426, + 953.4120559692383, + 953.1736373901367 + ], + "launch_us": [ + [ + 1014.4710540771484, + 9485.483169555664, + 6253.242492675781, + 960.5884552001953, + 9597.063064575195, + 10572.433471679688, + 6396.293640136719, + 9524.58381652832, + 7416.725158691406, + 6277.561187744141 + ], + [ + 953.4358978271484, + 954.1511535644531, + 952.4822235107422, + 953.9127349853516, + 954.3895721435547, + 952.9590606689453, + 952.2438049316406, + 954.3895721435547, + 952.7206420898438, + 953.4358978271484 + ], + [ + 952.7206420898438, + 953.4358978271484, + 952.9590606689453, + 953.1974792480469, + 952.4822235107422, + 952.9590606689453, + 953.1974792480469, + 954.3895721435547, + 953.1974792480469, + 953.1974792480469 + ] + ], + "median_us": 953.4120559692383, + "tflops": 360.38707663571455 + } + } + } + ] +} diff --git a/examples/ascend/flashmla/results/decode-correctness-final.json b/examples/ascend/flashmla/results/decode-correctness-final.json new file mode 100644 index 0000000000..c245426cec --- /dev/null +++ b/examples/ascend/flashmla/results/decode-correctness-final.json @@ -0,0 +1,116 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "id": 0, + "b": 4, + "sq": 1, + "topk": 128, + "extra_topk": null, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 1, + "b": 4, + "sq": 3, + "topk": 128, + "extra_topk": 128, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 2, + "b": 4, + "sq": 3, + "topk": 128, + "extra_topk": 128, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 3, + "b": 1, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 4, + "b": 74, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 5, + "b": 35, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 6, + "b": 1, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 7, + "b": 74, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + } + }, + { + "id": 8, + "b": 35, + "sq": 3, + "topk": 576, + "extra_topk": 64, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + } + } + ] +} diff --git a/examples/ascend/flashmla/results/decode-shared-fp4-recheck.json b/examples/ascend/flashmla/results/decode-shared-fp4-recheck.json new file mode 100644 index 0000000000..8b11ec855a --- /dev/null +++ b/examples/ascend/flashmla/results/decode-shared-fp4-recheck.json @@ -0,0 +1,168 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "id": 9, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 942.4686431884766, + 935.2684020996094, + 939.9652481079102, + 937.9386901855469, + 941.0619735717773 + ], + "launch_us": [ + [ + 930.3092956542969, + 949.859619140625, + 935.5545043945312, + 944.3759918212891, + 938.41552734375, + 945.5680847167969, + 943.1838989257812, + 944.1375732421875, + 940.3228759765625, + 952.9590606689453 + ], + [ + 937.7002716064453, + 935.0776672363281, + 936.7465972900391, + 933.6471557617188, + 935.3160858154297, + 933.4087371826172, + 938.1771087646484, + 932.2166442871094, + 935.0776672363281, + 935.3160858154297 + ], + [ + 929.8324584960938, + 949.859619140625, + 945.3296661376953, + 942.7070617675781, + 937.2234344482422, + 928.6403656005859, + 938.6539459228516, + 953.9127349853516, + 937.2234344482422, + 936.2697601318359 + ], + [ + 936.9850158691406, + 948.4291076660156, + 933.8855743408203, + 941.9918060302734, + 933.4087371826172, + 934.600830078125, + 934.1239929199219, + 933.4087371826172, + 945.5680847167969, + 936.9850158691406 + ], + [ + 936.7465972900391, + 935.7929229736328, + 935.0776672363281, + 938.8923645019531, + 936.7465972900391, + 954.1511535644531, + 934.8392486572266, + 951.0517120361328, + 954.3895721435547, + 932.9319000244141 + ] + ], + "median_us": 939.9652481079102, + "tflops": 365.54264572189186 + }, + "manual": { + "round_mean_us": [ + 952.2199630737305, + 953.0544281005859, + 952.4345397949219, + 953.1259536743164, + 952.6491165161133 + ], + "launch_us": [ + [ + 952.0053863525391, + 952.2438049316406, + 951.0517120361328, + 952.2438049316406, + 953.4358978271484, + 951.2901306152344, + 952.0053863525391, + 953.9127349853516, + 951.7669677734375, + 952.2438049316406 + ], + [ + 953.1974792480469, + 953.9127349853516, + 953.1974792480469, + 952.4822235107422, + 952.2438049316406, + 954.3895721435547, + 952.0053863525391, + 953.9127349853516, + 952.9590606689453, + 952.2438049316406 + ], + [ + 953.9127349853516, + 951.2901306152344, + 952.4822235107422, + 952.0053863525391, + 952.9590606689453, + 952.9590606689453, + 951.2901306152344, + 952.9590606689453, + 951.7669677734375, + 952.7206420898438 + ], + [ + 953.4358978271484, + 954.1511535644531, + 952.2438049316406, + 951.2901306152344, + 953.67431640625, + 953.4358978271484, + 952.9590606689453, + 953.4358978271484, + 952.2438049316406, + 954.3895721435547 + ], + [ + 954.3895721435547, + 951.2901306152344, + 952.7206420898438, + 953.67431640625, + 953.1974792480469, + 952.2438049316406, + 951.7669677734375, + 951.7669677734375, + 952.7206420898438, + 952.7206420898438 + ] + ], + "median_us": 952.6491165161133, + "tflops": 360.6756965634454 + } + } + } + ] +} diff --git a/examples/ascend/flashmla/results/environment.json b/examples/ascend/flashmla/results/environment.json new file mode 100644 index 0000000000..58621ab1c7 --- /dev/null +++ b/examples/ascend/flashmla/results/environment.json @@ -0,0 +1,51 @@ +{ + "tilelang_revision": "a35f8ddf45eba16c21211ec8822d56ce5363036f", + "reference_revision": "a123d0b0191e0da7aa1e044f8644e0989cc63220", + "python_version": "3.12.13", + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "cann": "/usr/local/Ascend/cann-9.2.0", + "bisheng": "2026-08-27T22:23:38+08:00 clang version 15.0.5 (clang-5c68a1cb1231 flang-5c68a1cb1231)\nTarget: aarch64-unknown-linux-gnu\nThread model: posix\nInstalledDir: /usr/local/Ascend/cann-9.2.0/bin\n", + "physical_device": 5, + "process_visible_device": 0, + "device": "Ascend950DT_9581", + "kernel_launch": { + "aic": 32, + "aiv_per_aic": 2, + "q_per_job": 1, + "heads": 64, + "dimension": 512, + "kv_block": 64 + }, + "pass_config": { + "tl.disable_shared_memory_reuse": true, + "tl.disable_safe_memory_legalize": true, + "tl.enable_auto_schedule": false + }, + "compile_flags": [ + "-Ofast", + "-mllvm", + "-enable-hiipu-vf-loop-unroll" + ], + "deterministic_algorithms": false, + "measurement_source_sha256": { + "examples/ascend/flashmla/decode.py": "8c156739d27ab2a0db7b68039c7eebac9e696aa1e3690e6c43029904399755f6", + "examples/ascend/flashmla/manual_intrinsics.h": "b1588c1bfeeccdd0fc4858cb5833bbd3a67a27a023df2292989678c8341ffe9c", + "examples/ascend/flashmla/quantization.py": "530dcc9133fb4d50b9ee278b8dee4e146ca016a20c67313e54e7f4b13c59efd0", + "examples/ascend/flashmla/core.py": "0a43fed18d416c498f43d44163037ac1a457b0f347aed7bcfff3a02b95d63e21", + "examples/ascend/flashmla/run.py": "01a9f4accc0302300c8ea8a6364120f2f4541eea20b99133d01123f2086e7a66", + "examples/ascend/flashmla/run_decode.py": "419de1ce91da6e0476abcb444d1ca0de684f25f5ee6388e9575bdfaee68cc567", + "examples/ascend/flashmla/manual.py": "0de7427f4514b84a58c4cf6f789dc0fba3f4f1617ebe05cebc9a08675c43abbb" + }, + "publication_note": "Hardware measurements were collected on the pinned tilelang_revision. The publication branch targets newer main; kernel and decode preparation sources are unchanged, while runners expose only reference and manual implementations.", + "publication_base_revision": "adeea32bcecdddca32195a95e90730b175606a0e", + "published_source_sha256": { + "examples/ascend/flashmla/core.py": "0a43fed18d416c498f43d44163037ac1a457b0f347aed7bcfff3a02b95d63e21", + "examples/ascend/flashmla/decode.py": "8c156739d27ab2a0db7b68039c7eebac9e696aa1e3690e6c43029904399755f6", + "examples/ascend/flashmla/manual.py": "0de7427f4514b84a58c4cf6f789dc0fba3f4f1617ebe05cebc9a08675c43abbb", + "examples/ascend/flashmla/manual_intrinsics.h": "b1588c1bfeeccdd0fc4858cb5833bbd3a67a27a023df2292989678c8341ffe9c", + "examples/ascend/flashmla/quantization.py": "530dcc9133fb4d50b9ee278b8dee4e146ca016a20c67313e54e7f4b13c59efd0", + "examples/ascend/flashmla/run.py": "5e673ca6266d6b7f090659b4004a4189d8103711fb6a1bafe6e38c2c50560247", + "examples/ascend/flashmla/run_decode.py": "a3c34e31a48e3e0cb61de567dab771e6027f5b06566811f860b5af84de467d5b" + } +} diff --git a/examples/ascend/flashmla/results/prefill-bench-final.json b/examples/ascend/flashmla/results/prefill-bench-final.json new file mode 100644 index 0000000000..56dc35ecbc --- /dev/null +++ b/examples/ascend/flashmla/results/prefill-bench-final.json @@ -0,0 +1,362 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "nq": 4096, + "nk": 4096, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 858.0684661865234, + 856.2564849853516, + 858.0446243286133 + ], + "launch_us": [ + [ + 874.2809295654297, + 868.7973022460938, + 849.2469787597656, + 856.8763732910156, + 848.0548858642578, + 860.6910705566406, + 859.7373962402344, + 861.4063262939453, + 853.7769317626953, + 847.8164672851562 + ], + [ + 855.6842803955078, + 850.9159088134766, + 854.2537689208984, + 858.7837219238281, + 869.9893951416016, + 849.4853973388672, + 849.9622344970703, + 849.4853973388672, + 875.4730224609375, + 848.5317230224609 + ], + [ + 851.8695831298828, + 868.0820465087891, + 850.2006530761719, + 850.2006530761719, + 861.8831634521484, + 848.5317230224609, + 852.3464202880859, + 852.8232574462891, + 869.2741394042969, + 875.2346038818359 + ] + ], + "median_us": 858.0446243286133, + "tflops": 400.44232425423286 + }, + "manual": { + "round_mean_us": [ + 848.6747741699219, + 849.3185043334961, + 848.388671875 + ], + "launch_us": [ + [ + 848.0548858642578, + 848.2933044433594, + 851.3927459716797, + 847.8164672851562, + 850.4390716552734, + 848.2933044433594, + 848.2933044433594, + 849.2469787597656, + 847.5780487060547, + 847.3396301269531 + ], + [ + 847.8164672851562, + 847.5780487060547, + 849.9622344970703, + 848.0548858642578, + 849.7238159179688, + 852.3464202880859, + 847.8164672851562, + 847.8164672851562, + 853.3000946044922, + 848.7701416015625 + ], + [ + 847.8164672851562, + 848.0548858642578, + 847.8164672851562, + 849.4853973388672, + 849.0085601806641, + 849.0085601806641, + 847.8164672851562, + 848.2933044433594, + 848.0548858642578, + 848.5317230224609 + ] + ], + "median_us": 848.6747741699219, + "tflops": 404.86343430682064 + } + } + }, + { + "nq": 4096, + "nk": 8192, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 853.7769317626953, + 856.9717407226562, + 851.6788482666016 + ], + "launch_us": [ + [ + 850.2006530761719, + 851.8695831298828, + 849.0085601806641, + 850.2006530761719, + 853.5385131835938, + 859.9758148193359, + 852.3464202880859, + 861.6447448730469, + 850.9159088134766, + 858.0684661865234 + ], + [ + 848.5317230224609, + 852.1080017089844, + 861.8831634521484, + 853.5385131835938, + 860.6910705566406, + 853.3000946044922, + 862.12158203125, + 851.6311645507812, + 860.2142333984375, + 865.6978607177734 + ], + [ + 847.5780487060547, + 850.2006530761719, + 849.4853973388672, + 851.6311645507812, + 850.2006530761719, + 856.1611175537109, + 851.1543273925781, + 851.8695831298828, + 852.1080017089844, + 856.3995361328125 + ] + ], + "median_us": 853.7769317626953, + "tflops": 402.4439767546939 + }, + "manual": { + "round_mean_us": [ + 850.7251739501953, + 851.1781692504883, + 851.1543273925781 + ], + "launch_us": [ + [ + 849.4853973388672, + 851.8695831298828, + 851.6311645507812, + 851.8695831298828, + 849.4853973388672, + 850.4390716552734, + 850.677490234375, + 849.4853973388672, + 851.3927459716797, + 850.9159088134766 + ], + [ + 850.4390716552734, + 850.677490234375, + 849.9622344970703, + 850.677490234375, + 850.4390716552734, + 851.8695831298828, + 850.677490234375, + 855.4458618164062, + 851.1543273925781, + 850.4390716552734 + ], + [ + 850.677490234375, + 850.9159088134766, + 850.677490234375, + 852.8232574462891, + 855.2074432373047, + 852.1080017089844, + 849.9622344970703, + 849.2469787597656, + 849.7238159179688, + 850.2006530761719 + ] + ], + "median_us": 851.1543273925781, + "tflops": 403.68400021248146 + } + } + }, + { + "nq": 4096, + "nk": 32768, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 861.5970611572266, + 862.2884750366211, + 862.1454238891602 + ], + "launch_us": [ + [ + 862.12158203125, + 862.3600006103516, + 861.1679077148438, + 861.4063262939453, + 862.12158203125, + 861.6447448730469, + 860.4526519775391, + 861.6447448730469, + 861.6447448730469, + 861.4063262939453 + ], + [ + 862.3600006103516, + 862.3600006103516, + 862.8368377685547, + 862.8368377685547, + 862.12158203125, + 861.6447448730469, + 861.6447448730469, + 862.3600006103516, + 862.8368377685547, + 861.8831634521484 + ], + [ + 862.8368377685547, + 861.8831634521484, + 862.3600006103516, + 861.6447448730469, + 862.12158203125, + 862.12158203125, + 862.3600006103516, + 862.3600006103516, + 862.12158203125, + 861.6447448730469 + ] + ], + "median_us": 862.1454238891602, + "tflops": 398.53761808538445 + }, + "manual": { + "round_mean_us": [ + 861.7877960205078, + 862.0500564575195, + 862.574577331543 + ], + "launch_us": [ + [ + 860.2142333984375, + 862.5984191894531, + 862.8368377685547, + 861.4063262939453, + 861.6447448730469, + 862.8368377685547, + 862.3600006103516, + 862.12158203125, + 860.4526519775391, + 861.4063262939453 + ], + [ + 863.0752563476562, + 859.9758148193359, + 862.5984191894531, + 863.0752563476562, + 862.8368377685547, + 862.8368377685547, + 861.1679077148438, + 859.4989776611328, + 862.5984191894531, + 862.8368377685547 + ], + [ + 862.8368377685547, + 863.3136749267578, + 862.3600006103516, + 862.12158203125, + 862.3600006103516, + 862.8368377685547, + 863.0752563476562, + 862.3600006103516, + 861.8831634521484, + 862.5984191894531 + ] + ], + "median_us": 862.0500564575195, + "tflops": 398.5817077629667 + } + } + } + ] +} diff --git a/examples/ascend/flashmla/results/prefill-code-size.json b/examples/ascend/flashmla/results/prefill-code-size.json new file mode 100644 index 0000000000..4faf800ed9 --- /dev/null +++ b/examples/ascend/flashmla/results/prefill-code-size.json @@ -0,0 +1,140 @@ +{ + "reference": { + "text_bytes": 14440, + "functions": [ + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE15clear_kv_buf_vfERU3AS6NS5_23UnifiedBufferMemoryPlanEj.vector.thread", + "bytes": 24 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE23kv_nd2nz_for_prefill_vfERU3AS6NS5_23UnifiedBufferMemoryPlanEj.vector.thread", + "bytes": 536 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE10softmax_vfILb0EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEffjjjjjj.vector.thread", + "bytes": 552 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE10softmax_vfILb1EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEffjjjjjj.vector.thread", + "bytes": 484 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE27get_final_lse_and_denorm_vfERU3AS6NS5_23UnifiedBufferMemoryPlanEjf.vector.thread", + "bytes": 96 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE19process_kv_index_vfIL9ModelType0EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEjjjij10FastDivMod.vector.thread", + "bytes": 184 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE18add_and_scale_o_vfILb1ELb0EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEjjjj.vector.thread", + "bytes": 52 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE18add_and_scale_o_vfILb1ELb1EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEjjjj.vector.thread", + "bytes": 68 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE19get_final_output_vfILb1EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEjjjjb.vector.thread", + "bytes": 136 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE19get_final_output_vfILb0EEEvRU3AS6NS5_23UnifiedBufferMemoryPlanEjjjjb.vector.thread", + "bytes": 100 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd22sparse_attn_fwd_kernelINS1_6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEEEEEvNT_6ParamsENS7_9AuxParamsE_mix_aic", + "bytes": 96 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE30sparse_attn_fwd_kernel_devfuncE19SparseAttnFwdParamsNS5_9AuxParamsE.cube", + "bytes": 4380 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd22sparse_attn_fwd_kernelINS1_6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEEEEEvNT_6ParamsENS7_9AuxParamsE_mix_aiv", + "bytes": 96 + }, + { + "symbol": "_ZN6ascend7prefill10sparse_fwd6KernelIXtlNS1_6ConfigEL17SparseAttnFwdMode0ELj64ELj640ELb1EEEE30sparse_attn_fwd_kernel_devfuncE19SparseAttnFwdParamsNS5_9AuxParamsE.vector", + "bytes": 6128 + } + ] + }, + "manual": { + "text_bytes": 15560, + "functions": [ + { + "symbol": "_ZL21main_kernel_simd_vf_0PU3AS6hi.vector.thread", + "bytes": 24 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_1PU3AS6hi.vector.thread", + "bytes": 536 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_3PU3AS6hiiii.vector.thread", + "bytes": 600 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_2PU3AS6hiiii.vector.thread", + "bytes": 552 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_4PU3AS6hi.vector.thread", + "bytes": 196 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_5PU3AS6hii.vector.thread", + "bytes": 168 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_6PU3AS6hiii.vector.thread", + "bytes": 124 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_7PU3AS6hiii.vector.thread", + "bytes": 100 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_8PU3AS6hii.vector.thread", + "bytes": 196 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_9PU3AS6hii.vector.thread", + "bytes": 140 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_8PU3AS6hii.vector.thread.1", + "bytes": 196 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_9PU3AS6hii.vector.thread.1", + "bytes": 140 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_8PU3AS6hii.vector.thread.2", + "bytes": 196 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_9PU3AS6hii.vector.thread.2", + "bytes": 140 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_8PU3AS6hii.vector.thread.3", + "bytes": 196 + }, + { + "symbol": "_ZL21main_kernel_simd_vf_9PU3AS6hii.vector.thread.3", + "bytes": 140 + }, + { + "symbol": "main_kernel_mix_aic", + "bytes": 3888 + }, + { + "symbol": "main_kernel_mix_aiv", + "bytes": 5880 + } + ] + } +} diff --git a/examples/ascend/flashmla/results/prefill-correctness-final.json b/examples/ascend/flashmla/results/prefill-correctness-final.json new file mode 100644 index 0000000000..5dd971ad95 --- /dev/null +++ b/examples/ascend/flashmla/results/prefill-correctness-final.json @@ -0,0 +1,330 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "nq": 1, + "nk": 128, + "topk": 128, + "sink": false, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 1, + "nk": 128, + "topk": 128, + "sink": true, + "lengths": true, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 1, + "nk": 128, + "topk": 128, + "sink": true, + "lengths": true, + "amplified": true, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 62, + "nk": 95, + "topk": 128, + "sink": false, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 62, + "nk": 95, + "topk": 128, + "sink": true, + "lengths": true, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 62, + "nk": 95, + "topk": 128, + "sink": true, + "lengths": true, + "amplified": true, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 512, + 256, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 213, + "nk": 1521, + "topk": 512, + "sink": false, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1280, + 640, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 213, + "nk": 1521, + "topk": 512, + "sink": true, + "lengths": true, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1280, + 640, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 213, + "nk": 1521, + "topk": 512, + "sink": true, + "lengths": true, + "amplified": true, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1280, + 640, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 65, + "nk": 4096, + "topk": 640, + "sink": false, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 65, + "nk": 4096, + "topk": 640, + "sink": true, + "lengths": true, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 65, + "nk": 4096, + "topk": 640, + "sink": true, + "lengths": true, + "amplified": true, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 65, + "nk": 128, + "topk": 64, + "sink": false, + "lengths": false, + "amplified": false, + "all_invalid_rows": true, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 384, + 192, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + }, + { + "nq": 33, + "nk": 95, + "topk": 192, + "sink": true, + "lengths": true, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 640, + 320, + 1 + ], + "correctness": { + "reference": true, + "manual": true + } + } + ] +} diff --git a/examples/ascend/flashmla/results/staged-b4-layouts-2026-10-20.json b/examples/ascend/flashmla/results/staged-b4-layouts-2026-10-20.json new file mode 100644 index 0000000000..3e902eda30 --- /dev/null +++ b/examples/ascend/flashmla/results/staged-b4-layouts-2026-10-20.json @@ -0,0 +1,343 @@ +{ + "date": "2026-10-20", + "source": "examples/ascend/flashmla/staged.py", + "source_sha256": "01b6b3b4a5703b447d8facd48a627f651cb4def9ec18c3315feea664d99f0727", + "baseline_commit": "9e691b088b49567afe0d4e4690a5e94983e75f84", + "baseline_source_sha256": "1cd8a2282ab2d960623fb85c4777c0b4d1679d158b9b6c87170e8fbafde55697", + "compiler_commit": "c536475fa3db8e41e04546388b6c76e15681190c", + "compiler_source_commits": [ + "44520d243cbd173e8aeaa12a15b52ef6501c3749", + "2a9f24798465fd93ebfe03682ffcd934b3e2dbf2" + ], + "native_library_sha256": "70a7dbb4d1e9eba8111985119c8d2960f0fc080c0604fbb2dcf8cbfc14858b80", + "target": "Ascend950DT_9581 / dav-3510", + "physical_device": 5, + "protocol": "Independent MajorK/MajorMN B4 layouts, default assume_conflict, shared physical ring clock, one S-to-A copy; two explicit PV GEMM alternatives.", + "l0b_bytes": 65536, + "validation": { + "bf16_correctness_cases": 14, + "decode_correctness_cases": 9, + "resident_bf16_tests": 6, + "resident_decode_boundary_tests": 4, + "compiler_host_tests": 741, + "compiler_device_regressions": 1 + }, + "rounds": 3, + "launches_per_round": 10, + "cases": [ + { + "nq": 4096, + "nk": 4096, + "topk": 640, + "performance": { + "best": { + "round_mean_us": [ + 865.8647537231445, + 865.483283996582, + 865.2925491333008 + ], + "median_us": 865.483283996582, + "launch_us": [ + [ + 864.9826049804688, + 865.4594421386719, + 865.4594421386719, + 865.6978607177734, + 865.2210235595703, + 865.2210235595703, + 865.6978607177734, + 867.1283721923828, + 867.6052093505859, + 866.1746978759766 + ], + [ + 867.3667907714844, + 864.2673492431641, + 864.9826049804688, + 864.7441864013672, + 864.9826049804688, + 864.0289306640625, + 867.6052093505859, + 865.4594421386719, + 864.9826049804688, + 866.4131164550781 + ], + [ + 865.6978607177734, + 864.0289306640625, + 864.5057678222656, + 865.6978607177734, + 866.1746978759766, + 867.3667907714844, + 864.9826049804688, + 865.2210235595703, + 864.0289306640625, + 865.2210235595703 + ] + ] + }, + "candidate": { + "round_mean_us": [ + 864.1481399536133, + 863.5997772216797, + 863.6474609375 + ], + "median_us": 863.6474609375, + "launch_us": [ + [ + 862.8368377685547, + 864.2673492431641, + 864.0289306640625, + 864.9826049804688, + 862.8368377685547, + 864.2673492431641, + 863.3136749267578, + 864.2673492431641, + 865.4594421386719, + 865.2210235595703 + ], + [ + 863.0752563476562, + 864.9826049804688, + 862.5984191894531, + 863.3136749267578, + 864.2673492431641, + 865.6978607177734, + 863.3136749267578, + 862.5984191894531, + 862.3600006103516, + 863.7905120849609 + ], + [ + 864.7441864013672, + 864.2673492431641, + 863.3136749267578, + 864.2673492431641, + 863.0752563476562, + 864.2673492431641, + 862.8368377685547, + 863.0752563476562, + 863.5520935058594, + 863.0752563476562 + ] + ] + } + }, + "candidate_change_percent": -0.2121153687226207, + "generated_source_sha256": { + "best": "392688682852a5e1adf4e387894302bfc8a95091e3c8a7e0c8be1bffc6cd10e6", + "candidate": "9d8767779a1ed0210be6f22342d8954e7c6103dffcad6f85a33c88f98a38b250" + } + }, + { + "nq": 4096, + "nk": 8192, + "topk": 640, + "performance": { + "best": { + "round_mean_us": [ + 864.0766143798828, + 864.5534515380859, + 864.4580841064453 + ], + "median_us": 864.4580841064453, + "launch_us": [ + [ + 865.2210235595703, + 864.2673492431641, + 863.5520935058594, + 864.2673492431641, + 863.5520935058594, + 863.5520935058594, + 864.7441864013672, + 864.0289306640625, + 863.7905120849609, + 863.7905120849609 + ], + [ + 863.7905120849609, + 864.7441864013672, + 864.2673492431641, + 862.5984191894531, + 869.5125579833984, + 863.7905120849609, + 864.9826049804688, + 863.0752563476562, + 864.2673492431641, + 864.5057678222656 + ], + [ + 864.2673492431641, + 864.9826049804688, + 864.9826049804688, + 864.2673492431641, + 866.4131164550781, + 863.0752563476562, + 866.4131164550781, + 864.2673492431641, + 862.8368377685547, + 863.0752563476562 + ] + ] + }, + "candidate": { + "round_mean_us": [ + 866.1508560180664, + 865.8170700073242, + 866.1270141601562 + ], + "median_us": 866.1270141601562, + "launch_us": [ + [ + 865.4594421386719, + 866.6515350341797, + 867.3667907714844, + 866.8899536132812, + 867.3667907714844, + 866.4131164550781, + 864.7441864013672, + 865.4594421386719, + 864.9826049804688, + 866.1746978759766 + ], + [ + 865.4594421386719, + 866.6515350341797, + 866.1746978759766, + 864.7441864013672, + 865.2210235595703, + 865.4594421386719, + 867.3667907714844, + 866.8899536132812, + 865.4594421386719, + 864.7441864013672 + ], + [ + 866.8899536132812, + 865.936279296875, + 865.936279296875, + 865.936279296875, + 867.3667907714844, + 864.9826049804688, + 865.6978607177734, + 865.936279296875, + 864.9826049804688, + 867.6052093505859 + ] + ] + } + }, + "candidate_change_percent": 0.19306084174526283, + "generated_source_sha256": { + "best": "8727ea3be148cea5fa537662c15b444ce3152525313d572ac5f03403a6fcfa6f", + "candidate": "667d927ebb75bb90d9bbbd5a40ca7a4d38c50589d7b612c10c68dc2a4a163ce2" + } + }, + { + "nq": 4096, + "nk": 32768, + "topk": 640, + "performance": { + "best": { + "round_mean_us": [ + 873.0173110961914, + 872.5643157958984, + 873.2318878173828 + ], + "median_us": 873.0173110961914, + "launch_us": [ + [ + 871.6583251953125, + 871.8967437744141, + 873.3272552490234, + 873.0888366699219, + 874.5193481445312, + 874.5193481445312, + 872.6119995117188, + 873.8040924072266, + 873.3272552490234, + 871.4199066162109 + ], + [ + 872.1351623535156, + 871.1814880371094, + 870.9430694580078, + 872.6119995117188, + 872.1351623535156, + 874.7577667236328, + 874.0425109863281, + 872.3735809326172, + 874.0425109863281, + 871.4199066162109 + ], + [ + 872.8504180908203, + 872.3735809326172, + 875.4730224609375, + 871.8967437744141, + 872.3735809326172, + 873.0888366699219, + 875.9498596191406, + 872.6119995117188, + 873.565673828125, + 872.1351623535156 + ] + ] + }, + "candidate": { + "round_mean_us": [ + 873.3987808227539, + 872.9219436645508, + 873.7564086914062 + ], + "median_us": 873.3987808227539, + "launch_us": [ + [ + 874.0425109863281, + 872.3735809326172, + 872.8504180908203, + 875.7114410400391, + 875.2346038818359, + 872.6119995117188, + 873.3272552490234, + 872.6119995117188, + 872.1351623535156, + 873.0888366699219 + ], + [ + 872.6119995117188, + 872.8504180908203, + 871.4199066162109, + 872.3735809326172, + 872.3735809326172, + 873.8040924072266, + 873.3272552490234, + 874.5193481445312, + 873.0888366699219, + 872.8504180908203 + ], + [ + 874.5193481445312, + 874.7577667236328, + 873.0888366699219, + 874.2809295654297, + 873.0888366699219, + 872.8504180908203, + 872.8504180908203, + 874.7577667236328, + 872.3735809326172, + 874.9961853027344 + ] + ] + } + }, + "candidate_change_percent": 0.043695551246680786, + "generated_source_sha256": { + "best": "c43075d696174ca6beff32f61217dbfb424314c7101a7e23d62fe2036dccabcf", + "candidate": "4b8bdb1e6902ec234abe2f5e4383895e5e1d1a446e75db5d0521a0c6cdd5914d" + } + } + ], + "evidence_directory": ".tlcli/flashmla-phase-perf/b4-assume-conflict-integration-20261020", + "decode_performance_sweep": "Not remeasured in this change; historical results identify their own source and compiler." +} diff --git a/examples/ascend/flashmla/results/staged-decode-bench-2026-10-18.json b/examples/ascend/flashmla/results/staged-decode-bench-2026-10-18.json new file mode 100644 index 0000000000..3db82c5d87 --- /dev/null +++ b/examples/ascend/flashmla/results/staged-decode-bench-2026-10-18.json @@ -0,0 +1,1589 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "5", + "cases": [ + { + "id": 0, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 78.4158706665039, + 78.53507995605469, + 78.48739624023438 + ], + "launch_us": [ + [ + 78.67813110351562, + 78.91654968261719, + 77.72445678710938, + 78.91654968261719, + 78.67813110351562, + 78.43971252441406, + 77.72445678710938, + 78.43971252441406, + 78.2012939453125, + 78.43971252441406 + ], + [ + 78.67813110351562, + 79.15496826171875, + 78.2012939453125, + 79.87022399902344, + 78.43971252441406, + 78.43971252441406, + 78.67813110351562, + 77.24761962890625, + 78.67813110351562, + 77.96287536621094 + ], + [ + 79.15496826171875, + 78.2012939453125, + 78.43971252441406, + 77.96287536621094, + 78.2012939453125, + 78.91654968261719, + 79.15496826171875, + 78.91654968261719, + 76.77078247070312, + 79.15496826171875 + ] + ], + "median_us": 78.48739624023438, + "tflops": 273.60872584267895 + }, + "manual": { + "round_mean_us": [ + 76.48468017578125, + 76.81846618652344, + 76.55620574951172 + ], + "launch_us": [ + [ + 76.05552673339844, + 75.81710815429688, + 77.48603820800781, + 76.2939453125, + 76.53236389160156, + 76.2939453125, + 75.81710815429688, + 76.53236389160156, + 76.77078247070312, + 77.24761962890625 + ], + [ + 76.53236389160156, + 77.48603820800781, + 76.05552673339844, + 76.53236389160156, + 77.00920104980469, + 77.96287536621094, + 77.72445678710938, + 76.2939453125, + 76.05552673339844, + 76.53236389160156 + ], + [ + 76.2939453125, + 76.05552673339844, + 76.05552673339844, + 76.2939453125, + 77.00920104980469, + 77.96287536621094, + 77.00920104980469, + 75.34027099609375, + 77.00920104980469, + 76.53236389160156 + ] + ], + "median_us": 76.55620574951172, + "tflops": 280.51072110685124 + }, + "staged": { + "round_mean_us": [ + 71.97856903076172, + 71.21562957763672, + 72.45540618896484 + ], + "launch_us": [ + [ + 72.479248046875, + 71.52557373046875, + 72.00241088867188, + 71.04873657226562, + 72.24082946777344, + 72.479248046875, + 72.479248046875, + 71.52557373046875, + 71.76399230957031, + 72.24082946777344 + ], + [ + 70.5718994140625, + 71.04873657226562, + 71.28715515136719, + 72.24082946777344, + 71.04873657226562, + 71.76399230957031, + 72.00241088867188, + 70.5718994140625, + 71.52557373046875, + 70.09506225585938 + ], + [ + 72.71766662597656, + 73.19450378417969, + 72.24082946777344, + 72.00241088867188, + 73.19450378417969, + 72.00241088867188, + 73.67134094238281, + 72.00241088867188, + 71.76399230957031, + 71.76399230957031 + ] + ], + "median_us": 71.97856903076172, + "tflops": 298.3504224823118 + } + }, + "evidence_source": "matrix-v91-decode0.json" + }, + { + "id": 1, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 144.7916030883789, + 144.3624496459961, + 144.43397521972656 + ], + "launch_us": [ + [ + 146.1505889892578, + 144.72007751464844, + 146.62742614746094, + 143.52798461914062, + 145.19691467285156, + 144.72007751464844, + 143.28956604003906, + 144.00482177734375, + 144.95849609375, + 144.72007751464844 + ], + [ + 145.91217041015625, + 143.52798461914062, + 144.2432403564453, + 143.52798461914062, + 144.48165893554688, + 144.2432403564453, + 143.7664031982422, + 145.43533325195312, + 144.2432403564453, + 144.2432403564453 + ], + [ + 144.00482177734375, + 145.19691467285156, + 143.0511474609375, + 145.91217041015625, + 143.52798461914062, + 145.43533325195312, + 143.52798461914062, + 144.2432403564453, + 146.1505889892578, + 143.28956604003906 + ] + ], + "median_us": 144.43397521972656, + "tflops": 297.3654425467479 + }, + "manual": { + "round_mean_us": [ + 143.02730560302734, + 142.93193817138672, + 143.24188232421875 + ], + "launch_us": [ + [ + 143.28956604003906, + 142.3358917236328, + 141.8590545654297, + 143.28956604003906, + 143.7664031982422, + 144.00482177734375, + 144.72007751464844, + 141.143798828125, + 141.143798828125, + 144.72007751464844 + ], + [ + 143.28956604003906, + 142.81272888183594, + 142.81272888183594, + 142.57431030273438, + 142.81272888183594, + 142.57431030273438, + 143.7664031982422, + 144.00482177734375, + 142.3358917236328, + 142.3358917236328 + ], + [ + 144.00482177734375, + 142.81272888183594, + 142.57431030273438, + 143.52798461914062, + 143.52798461914062, + 144.00482177734375, + 142.3358917236328, + 141.8590545654297, + 143.28956604003906, + 144.48165893554688 + ] + ], + "median_us": 143.02730560302734, + "tflops": 300.2900234952823 + }, + "staged": { + "round_mean_us": [ + 136.06548309326172, + 136.30390167236328, + 136.04164123535156 + ], + "launch_us": [ + [ + 137.09068298339844, + 134.94491577148438, + 136.85226440429688, + 135.89859008789062, + 136.6138458251953, + 136.85226440429688, + 134.7064971923828, + 135.66017150878906, + 135.4217529296875, + 136.6138458251953 + ], + [ + 134.7064971923828, + 136.6138458251953, + 135.89859008789062, + 137.80593872070312, + 137.09068298339844, + 135.89859008789062, + 137.3291015625, + 135.89859008789062, + 135.18333435058594, + 136.6138458251953 + ], + [ + 135.4217529296875, + 137.09068298339844, + 137.80593872070312, + 134.94491577148438, + 135.4217529296875, + 136.1370086669922, + 136.6138458251953, + 136.37542724609375, + 134.94491577148438, + 135.66017150878906 + ] + ], + "median_us": 136.06548309326172, + "tflops": 315.6544333184157 + } + }, + "evidence_source": "matrix-v91-decode1.json" + }, + { + "id": 2, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 276.97086334228516, + 276.87549591064453, + 277.3761749267578 + ], + "launch_us": [ + [ + 276.32713317871094, + 276.80397033691406, + 278.71131896972656, + 277.0423889160156, + 276.0887145996094, + 275.3734588623047, + 277.2808074951172, + 277.51922607421875, + 276.0887145996094, + 278.472900390625 + ], + [ + 275.8502960205078, + 275.3734588623047, + 276.80397033691406, + 275.61187744140625, + 279.1881561279297, + 277.2808074951172, + 275.8502960205078, + 277.2808074951172, + 277.9960632324219, + 277.51922607421875 + ], + [ + 277.7576446533203, + 278.472900390625, + 275.8502960205078, + 277.2808074951172, + 277.51922607421875, + 278.23448181152344, + 277.9960632324219, + 276.0887145996094, + 277.0423889160156, + 277.51922607421875 + ] + ], + "median_us": 276.97086334228516, + "tflops": 310.1385643364377 + }, + "manual": { + "round_mean_us": [ + 272.6316452026367, + 273.27537536621094, + 273.5137939453125 + ], + "launch_us": [ + [ + 273.2276916503906, + 271.5587615966797, + 273.70452880859375, + 272.98927307128906, + 271.3203430175781, + 272.0355987548828, + 272.7508544921875, + 273.4661102294922, + 272.51243591308594, + 272.7508544921875 + ], + [ + 275.61187744140625, + 274.89662170410156, + 272.2740173339844, + 272.2740173339844, + 273.2276916503906, + 272.2740173339844, + 273.9429473876953, + 273.2276916503906, + 272.2740173339844, + 272.7508544921875 + ], + [ + 272.7508544921875, + 274.658203125, + 273.2276916503906, + 273.4661102294922, + 272.7508544921875, + 273.70452880859375, + 272.98927307128906, + 274.89662170410156, + 273.9429473876953, + 272.7508544921875 + ] + ], + "median_us": 273.27537536621094, + "tflops": 314.332551203664 + }, + "staged": { + "round_mean_us": [ + 259.4470977783203, + 260.162353515625, + 260.30540466308594 + ], + "launch_us": [ + [ + 259.3994140625, + 258.44573974609375, + 260.5915069580078, + 258.2073211669922, + 258.44573974609375, + 258.9225769042969, + 258.6841583251953, + 260.5915069580078, + 261.54518127441406, + 259.63783264160156 + ], + [ + 260.1146697998047, + 262.0220184326172, + 264.4062042236328, + 259.16099548339844, + 258.6841583251953, + 261.06834411621094, + 259.63783264160156, + 257.9689025878906, + 258.9225769042969, + 259.63783264160156 + ], + [ + 258.44573974609375, + 258.9225769042969, + 262.97569274902344, + 261.54518127441406, + 260.35308837890625, + 260.5915069580078, + 257.4920654296875, + 260.1146697998047, + 260.8299255371094, + 261.7835998535156 + ] + ], + "median_us": 260.162353515625, + "tflops": 330.1759257603003 + } + }, + "evidence_source": "matrix-v91-decode2.json" + }, + { + "id": 3, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 546.2884902954102, + 547.2898483276367, + 546.2646484375 + ], + "launch_us": [ + [ + 546.6938018798828, + 549.31640625, + 545.2632904052734, + 544.7864532470703, + 545.9785461425781, + 545.2632904052734, + 547.1706390380859, + 546.2169647216797, + 548.3627319335938, + 543.8327789306641 + ], + [ + 548.1243133544922, + 546.2169647216797, + 546.2169647216797, + 548.6011505126953, + 546.6938018798828, + 547.4090576171875, + 546.2169647216797, + 550.2700805664062, + 546.9322204589844, + 546.2169647216797 + ], + [ + 545.2632904052734, + 548.1243133544922, + 545.501708984375, + 546.6938018798828, + 546.9322204589844, + 543.8327789306641, + 546.9322204589844, + 546.9322204589844, + 546.4553833007812, + 545.9785461425781 + ] + ], + "median_us": 546.2884902954102, + "tflops": 314.4834549728448 + }, + "manual": { + "round_mean_us": [ + 541.3293838500977, + 540.5426025390625, + 540.7094955444336 + ], + "launch_us": [ + [ + 542.6406860351562, + 544.7864532470703, + 537.1570587158203, + 539.7796630859375, + 541.2101745605469, + 542.8791046142578, + 539.5412445068359, + 544.3096160888672, + 540.2565002441406, + 540.7333374023438 + ], + [ + 540.7333374023438, + 538.1107330322266, + 539.7796630859375, + 540.2565002441406, + 540.0180816650391, + 539.0644073486328, + 545.9785461425781, + 539.7796630859375, + 540.7333374023438, + 540.9717559814453 + ], + [ + 540.2565002441406, + 539.7796630859375, + 541.68701171875, + 540.9717559814453, + 540.2565002441406, + 541.68701171875, + 543.1175231933594, + 540.7333374023438, + 538.1107330322266, + 540.4949188232422 + ] + ], + "median_us": 540.7094955444336, + "tflops": 317.7282686094093 + }, + "staged": { + "round_mean_us": [ + 511.02638244628906, + 511.31248474121094, + 512.5284194946289 + ], + "launch_us": [ + [ + 511.16943359375, + 511.6462707519531, + 510.6925964355469, + 509.7389221191406, + 513.7920379638672, + 512.3615264892578, + 509.9773406982422, + 510.6925964355469, + 509.9773406982422, + 510.21575927734375 + ], + [ + 510.93101501464844, + 508.7852478027344, + 513.0767822265625, + 511.6462707519531, + 510.93101501464844, + 511.6462707519531, + 512.3615264892578, + 509.7389221191406, + 513.0767822265625, + 510.93101501464844 + ], + [ + 513.7920379638672, + 511.8846893310547, + 513.0767822265625, + 512.8383636474609, + 512.3615264892578, + 511.6462707519531, + 512.5999450683594, + 515.4609680175781, + 511.40785217285156, + 510.21575927734375 + ] + ], + "median_us": 511.31248474121094, + "tflops": 335.9954958403802 + } + }, + "evidence_source": "matrix-v91-decode3.json" + }, + { + "id": 4, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 70.45269012451172, + 70.85800170898438, + 71.21562957763672 + ], + "launch_us": [ + [ + 69.61822509765625, + 69.37980651855469, + 71.52557373046875, + 69.85664367675781, + 69.14138793945312, + 71.28715515136719, + 70.81031799316406, + 70.5718994140625, + 69.85664367675781, + 72.479248046875 + ], + [ + 68.66455078125, + 71.28715515136719, + 69.61822509765625, + 71.28715515136719, + 72.24082946777344, + 71.28715515136719, + 72.71766662597656, + 72.71766662597656, + 68.90296936035156, + 69.85664367675781 + ], + [ + 68.90296936035156, + 72.479248046875, + 72.479248046875, + 72.00241088867188, + 70.81031799316406, + 71.04873657226562, + 71.52557373046875, + 69.37980651855469, + 71.76399230957031, + 71.76399230957031 + ] + ], + "median_us": 70.85800170898438, + "tflops": 303.0686155700199 + }, + "manual": { + "round_mean_us": [ + 70.30963897705078, + 70.14274597167969, + 70.11890411376953 + ], + "launch_us": [ + [ + 69.85664367675781, + 72.479248046875, + 70.81031799316406, + 68.18771362304688, + 71.76399230957031, + 68.66455078125, + 69.85664367675781, + 68.42613220214844, + 71.28715515136719, + 71.76399230957031 + ], + [ + 68.90296936035156, + 68.66455078125, + 70.81031799316406, + 69.14138793945312, + 72.24082946777344, + 69.37980651855469, + 71.28715515136719, + 69.85664367675781, + 71.28715515136719, + 69.85664367675781 + ], + [ + 69.14138793945312, + 71.28715515136719, + 69.61822509765625, + 67.94929504394531, + 70.33348083496094, + 72.00241088867188, + 69.14138793945312, + 71.28715515136719, + 70.5718994140625, + 69.85664367675781 + ] + ], + "median_us": 70.14274597167969, + "tflops": 306.1590501271581 + }, + "staged": { + "round_mean_us": [ + 69.2605972290039, + 68.75991821289062, + 68.95065307617188 + ], + "launch_us": [ + [ + 70.5718994140625, + 68.42613220214844, + 67.94929504394531, + 68.90296936035156, + 71.52557373046875, + 67.71087646484375, + 68.90296936035156, + 68.90296936035156, + 69.61822509765625, + 70.09506225585938 + ], + [ + 70.09506225585938, + 69.61822509765625, + 68.18771362304688, + 67.94929504394531, + 68.42613220214844, + 67.94929504394531, + 68.90296936035156, + 67.94929504394531, + 68.66455078125, + 69.85664367675781 + ], + [ + 68.42613220214844, + 68.18771362304688, + 68.66455078125, + 70.33348083496094, + 69.14138793945312, + 68.42613220214844, + 70.09506225585938, + 69.37980651855469, + 68.66455078125, + 68.18771362304688 + ] + ], + "median_us": 68.95065307617188, + "tflops": 311.4522563880011 + } + }, + "evidence_source": "matrix-v100-decode4.json" + }, + { + "id": 5, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 130.22422790527344, + 131.48784637451172, + 130.74874877929688 + ], + "launch_us": [ + [ + 128.5076141357422, + 131.13021850585938, + 130.8917999267578, + 131.6070556640625, + 129.69970703125, + 129.69970703125, + 129.22286987304688, + 131.6070556640625, + 128.5076141357422, + 131.36863708496094 + ], + [ + 132.3223114013672, + 132.3223114013672, + 132.56072998046875, + 132.56072998046875, + 129.46128845214844, + 131.13021850585938, + 131.84547424316406, + 131.36863708496094, + 132.3223114013672, + 128.9844512939453 + ], + [ + 132.7991485595703, + 131.36863708496094, + 130.4149627685547, + 130.17654418945312, + 130.8917999267578, + 131.13021850585938, + 130.8917999267578, + 129.46128845214844, + 131.13021850585938, + 129.22286987304688 + ] + ], + "median_us": 130.74874877929688, + "tflops": 328.4901259934716 + }, + "manual": { + "round_mean_us": [ + 131.0110092163086, + 131.79779052734375, + 131.15406036376953 + ], + "launch_us": [ + [ + 131.84547424316406, + 130.17654418945312, + 129.93812561035156, + 132.3223114013672, + 131.84547424316406, + 129.93812561035156, + 130.17654418945312, + 132.08389282226562, + 130.4149627685547, + 131.36863708496094 + ], + [ + 130.4149627685547, + 136.6138458251953, + 132.08389282226562, + 132.56072998046875, + 130.8917999267578, + 129.93812561035156, + 129.93812561035156, + 132.08389282226562, + 131.84547424316406, + 131.6070556640625 + ], + [ + 132.3223114013672, + 130.65338134765625, + 132.08389282226562, + 130.8917999267578, + 131.13021850585938, + 131.13021850585938, + 132.08389282226562, + 129.46128845214844, + 130.4149627685547, + 131.36863708496094 + ] + ], + "median_us": 131.15406036376953, + "tflops": 327.4749774492271 + }, + "staged": { + "round_mean_us": [ + 129.1513442993164, + 129.00829315185547, + 128.43608856201172 + ], + "launch_us": [ + [ + 128.9844512939453, + 127.55393981933594, + 130.8917999267578, + 129.69970703125, + 129.46128845214844, + 128.9844512939453, + 129.46128845214844, + 129.46128845214844, + 129.22286987304688, + 127.7923583984375 + ], + [ + 128.74603271484375, + 128.74603271484375, + 129.46128845214844, + 128.03077697753906, + 128.26919555664062, + 129.69970703125, + 128.03077697753906, + 128.74603271484375, + 130.17654418945312, + 130.17654418945312 + ], + [ + 128.74603271484375, + 130.17654418945312, + 128.74603271484375, + 126.60026550292969, + 126.60026550292969, + 129.46128845214844, + 127.7923583984375, + 128.5076141357422, + 129.46128845214844, + 128.26919555664062 + ] + ], + "median_us": 129.00829315185547, + "tflops": 332.92179836410986 + } + }, + "evidence_source": "matrix-v100-decode5.json" + }, + { + "id": 6, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 250.17261505126953, + 251.0547637939453, + 251.38854980468753 + ], + "launch_us": [ + [ + 249.8626708984375, + 250.81634521484375, + 250.10108947753906, + 250.81634521484375, + 250.10108947753906, + 248.19374084472656, + 250.5779266357422, + 250.33950805664062, + 250.81634521484375, + 250.10108947753906 + ], + [ + 251.0547637939453, + 251.53160095214844, + 249.8626708984375, + 249.8626708984375, + 251.29318237304688, + 248.43215942382812, + 252.72369384765625, + 252.72369384765625, + 248.90899658203125, + 254.15420532226562 + ], + [ + 252.72369384765625, + 251.29318237304688, + 251.0547637939453, + 250.5779266357422, + 250.10108947753906, + 253.91578674316406, + 248.6705780029297, + 250.10108947753906, + 251.77001953125, + 253.6773681640625 + ] + ], + "median_us": 251.0547637939453, + "tflops": 342.1538178439123 + }, + "manual": { + "round_mean_us": [ + 252.50911712646482, + 252.43759155273438, + 252.72369384765625 + ], + "launch_us": [ + [ + 251.53160095214844, + 253.91578674316406, + 252.72369384765625, + 251.53160095214844, + 253.43894958496094, + 252.00843811035156, + 253.43894958496094, + 251.53160095214844, + 252.72369384765625, + 252.24685668945312 + ], + [ + 253.20053100585938, + 251.77001953125, + 252.72369384765625, + 252.72369384765625, + 251.53160095214844, + 252.9621124267578, + 252.00843811035156, + 254.63104248046875, + 250.5779266357422, + 252.24685668945312 + ], + [ + 254.3926239013672, + 252.24685668945312, + 250.81634521484375, + 253.43894958496094, + 253.43894958496094, + 252.24685668945312, + 252.9621124267578, + 252.00843811035156, + 252.4852752685547, + 253.20053100585938 + ] + ], + "median_us": 252.50911712646482, + "tflops": 340.18314624647314 + }, + "staged": { + "round_mean_us": [ + 246.95396423339844, + 245.9049224853516, + 246.3102340698242 + ], + "launch_us": [ + [ + 246.52481079101562, + 247.00164794921875, + 247.00164794921875, + 247.00164794921875, + 247.2400665283203, + 246.7632293701172, + 246.7632293701172, + 246.7632293701172, + 246.52481079101562, + 247.955322265625 + ], + [ + 247.00164794921875, + 245.3327178955078, + 247.00164794921875, + 244.61746215820312, + 244.8558807373047, + 246.28639221191406, + 245.09429931640625, + 246.0479736328125, + 245.80955505371094, + 247.00164794921875 + ], + [ + 245.80955505371094, + 245.80955505371094, + 244.61746215820312, + 248.19374084472656, + 247.00164794921875, + 247.00164794921875, + 245.80955505371094, + 246.7632293701172, + 245.80955505371094, + 246.28639221191406 + ] + ], + "median_us": 246.3102340698242, + "tflops": 348.7445263668955 + } + }, + "evidence_source": "matrix-v100-decode6.json" + }, + { + "id": 7, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 496.05369567871094, + 497.07889556884766, + 497.0312118530273 + ], + "launch_us": [ + [ + 496.3874816894531, + 494.7185516357422, + 497.1027374267578, + 494.0032958984375, + 498.5332489013672, + 495.91064453125, + 495.91064453125, + 494.95697021484375, + 495.91064453125, + 497.1027374267578 + ], + [ + 498.05641174316406, + 494.7185516357422, + 493.04962158203125, + 498.2948303222656, + 500.67901611328125, + 497.57957458496094, + 497.57957458496094, + 497.1027374267578, + 497.3411560058594, + 496.3874816894531 + ], + [ + 495.91064453125, + 494.95697021484375, + 497.1027374267578, + 495.91064453125, + 502.1095275878906, + 497.3411560058594, + 496.6259002685547, + 498.5332489013672, + 496.3874816894531, + 495.4338073730469 + ] + ], + "median_us": 497.0312118530273, + "tflops": 345.649705175459 + }, + "manual": { + "round_mean_us": [ + 500.17833709716797, + 500.631332397461, + 501.8472671508789 + ], + "launch_us": [ + [ + 500.67901611328125, + 504.49371337890625, + 498.77166748046875, + 500.67901611328125, + 497.57957458496094, + 499.0100860595703, + 497.8179931640625, + 500.2021789550781, + 499.2485046386719, + 503.30162048339844 + ], + [ + 499.0100860595703, + 502.8247833251953, + 501.1558532714844, + 497.8179931640625, + 500.4405975341797, + 501.1558532714844, + 500.2021789550781, + 501.39427185058594, + 499.725341796875, + 502.58636474609375 + ], + [ + 499.725341796875, + 503.5400390625, + 506.1626434326172, + 501.87110900878906, + 504.2552947998047, + 500.67901611328125, + 500.4405975341797, + 500.67901611328125, + 502.58636474609375, + 498.5332489013672 + ] + ], + "median_us": 500.631332397461, + "tflops": 343.16408247417814 + }, + "staged": { + "round_mean_us": [ + 479.14981842041016, + 479.6743392944336, + 480.3180694580078 + ], + "launch_us": [ + [ + 482.0823669433594, + 479.2213439941406, + 478.98292541503906, + 478.98292541503906, + 478.7445068359375, + 478.2676696777344, + 478.0292510986328, + 480.1750183105469, + 478.0292510986328, + 478.98292541503906 + ], + [ + 477.79083251953125, + 478.7445068359375, + 479.9365997314453, + 481.3671112060547, + 478.98292541503906, + 478.7445068359375, + 479.2213439941406, + 480.89027404785156, + 480.1750183105469, + 480.89027404785156 + ], + [ + 481.8439483642578, + 483.0360412597656, + 478.50608825683594, + 478.50608825683594, + 479.4597625732422, + 479.2213439941406, + 482.0823669433594, + 478.98292541503906, + 482.5592041015625, + 478.98292541503906 + ] + ], + "median_us": 479.6743392944336, + "tflops": 358.15693641795286 + } + }, + "evidence_source": "matrix-v100-decode7.json" + }, + { + "id": 8, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 1045.823097229004, + 1045.4893112182617, + 1045.7038879394531 + ], + "launch_us": [ + [ + 1044.7502136230469, + 1045.9423065185547, + 1046.1807250976562, + 1046.6575622558594, + 1045.9423065185547, + 1047.372817993164, + 1044.7502136230469, + 1045.22705078125, + 1044.5117950439453, + 1046.895980834961 + ], + [ + 1042.6044464111328, + 1047.1343994140625, + 1043.7965393066406, + 1044.9886322021484, + 1044.5117950439453, + 1046.895980834961, + 1045.9423065185547, + 1046.6575622558594, + 1045.4654693603516, + 1046.895980834961 + ], + [ + 1045.9423065185547, + 1045.22705078125, + 1045.9423065185547, + 1046.1807250976562, + 1046.6575622558594, + 1045.9423065185547, + 1044.2733764648438, + 1044.9886322021484, + 1046.1807250976562, + 1045.7038879394531 + ] + ], + "median_us": 1045.7038879394531, + "tflops": 328.58000017295 + }, + "manual": { + "round_mean_us": [ + 1034.4743728637695, + 1034.6651077270508, + 1034.235954284668 + ], + "launch_us": [ + [ + 1034.0213775634766, + 1034.2597961425781, + 1036.6439819335938, + 1033.0677032470703, + 1036.1671447753906, + 1035.928726196289, + 1032.5908660888672, + 1034.0213775634766, + 1033.5445404052734, + 1034.4982147216797 + ], + [ + 1034.2597961425781, + 1036.6439819335938, + 1032.5908660888672, + 1034.2597961425781, + 1036.4055633544922, + 1033.5445404052734, + 1035.6903076171875, + 1032.3524475097656, + 1035.451889038086, + 1035.451889038086 + ], + [ + 1035.6903076171875, + 1034.0213775634766, + 1032.8292846679688, + 1034.9750518798828, + 1035.6903076171875, + 1032.5908660888672, + 1035.2134704589844, + 1033.782958984375, + 1033.782958984375, + 1033.782958984375 + ] + ], + "median_us": 1034.4743728637695, + "tflops": 332.14683001649234 + }, + "staged": { + "round_mean_us": [ + 973.7014770507812, + 974.3452072143555, + 974.1783142089844 + ], + "launch_us": [ + [ + 974.8935699462891, + 973.9398956298828, + 974.4167327880859, + 972.747802734375, + 974.8935699462891, + 973.2246398925781, + 972.9862213134766, + 972.747802734375, + 973.7014770507812, + 973.4630584716797 + ], + [ + 973.9398956298828, + 975.3704071044922, + 974.6551513671875, + 974.8935699462891, + 972.9862213134766, + 975.1319885253906, + 974.6551513671875, + 973.7014770507812, + 974.4167327880859, + 973.7014770507812 + ], + [ + 975.3704071044922, + 973.4630584716797, + 974.6551513671875, + 972.9862213134766, + 973.4630584716797, + 974.6551513671875, + 975.1319885253906, + 974.8935699462891, + 973.2246398925781, + 973.9398956298828 + ] + ], + "median_us": 974.1783142089844, + "tflops": 352.70481663205055 + } + }, + "evidence_source": "matrix-v91-decode8.json" + }, + { + "id": 9, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 941.0619735717773, + 937.6287460327148, + 936.2459182739258 + ], + "launch_us": [ + [ + 934.3624114990234, + 936.5081787109375, + 936.0313415527344, + 933.6471557617188, + 956.0585021972656, + 935.0776672363281, + 947.7138519287109, + 934.3624114990234, + 943.6607360839844, + 953.1974792480469 + ], + [ + 933.8855743408203, + 945.8065032958984, + 936.2697601318359, + 934.3624114990234, + 935.0776672363281, + 930.5477142333984, + 935.5545043945312, + 934.8392486572266, + 952.4822235107422, + 937.4618530273438 + ], + [ + 933.8855743408203, + 948.6675262451172, + 934.1239929199219, + 930.3092956542969, + 931.0245513916016, + 934.600830078125, + 934.600830078125, + 934.600830078125, + 937.4618530273438, + 943.1838989257812 + ] + ], + "median_us": 937.6287460327148, + "tflops": 366.4535511883842 + }, + "manual": { + "round_mean_us": [ + 953.8412094116211, + 954.8187255859375, + 954.6041488647461 + ], + "launch_us": [ + [ + 953.1974792480469, + 954.6279907226562, + 953.4358978271484, + 954.3895721435547, + 953.4358978271484, + 953.1974792480469, + 954.3895721435547, + 954.3895721435547, + 953.9127349853516, + 953.4358978271484 + ], + [ + 956.0585021972656, + 956.0585021972656, + 954.3895721435547, + 954.8664093017578, + 953.9127349853516, + 955.1048278808594, + 954.8664093017578, + 953.4358978271484, + 953.9127349853516, + 955.5816650390625 + ], + [ + 954.3895721435547, + 955.5816650390625, + 955.8200836181641, + 953.67431640625, + 954.3895721435547, + 954.1511535644531, + 954.3895721435547, + 955.5816650390625, + 954.6279907226562, + 953.4358978271484 + ] + ], + "median_us": 954.6041488647461, + "tflops": 359.93703158384545 + }, + "staged": { + "round_mean_us": [ + 935.7452392578125, + 935.6260299682617, + 935.7690811157227 + ], + "launch_us": [ + [ + 935.3160858154297, + 935.0776672363281, + 936.0313415527344, + 935.7929229736328, + 934.3624114990234, + 940.3228759765625, + 935.3160858154297, + 935.0776672363281, + 935.7929229736328, + 934.3624114990234 + ], + [ + 934.600830078125, + 937.4618530273438, + 935.3160858154297, + 936.0313415527344, + 935.5545043945312, + 934.8392486572266, + 935.0776672363281, + 935.5545043945312, + 936.5081787109375, + 935.3160858154297 + ], + [ + 936.0313415527344, + 936.0313415527344, + 936.2697601318359, + 935.7929229736328, + 935.5545043945312, + 937.2234344482422, + 934.8392486572266, + 935.5545043945312, + 935.3160858154297, + 935.0776672363281 + ] + ], + "median_us": 935.7452392578125, + "tflops": 367.1911640742353 + } + }, + "evidence_source": "matrix-v100-decode9.json" + } + ], + "provenance": { + "candidate_source_sha256": "459c1a0133caab9b182172f187313d288509ba7d94596239c4df6b110a4bc296", + "native_sha256": "b45698ea1ba4a96a71a782f979d1ec88c332b19c3fb278dd77a175af6fad74bf", + "unchanged_FP8_source_equivalence": "v100-shared-raw-task-variant/UNCHANGED_PATH_PROOF.json", + "status": "All 10 performance, 9 correctness, 4 independent-length boundary cases, and affected permanent pytest nodes PASS", + "compiler_commit": "0d86dd0b886624bb2e35bbfdf83e9e30cc358c98", + "source_sha256": "f93885c71bba1103c1b98569a1cad4c5a56d044e7d47207f5b502ca26ed1aa4e", + "measurement_date": "2026-10-18", + "reference_commit": "a123d0b0191e0da7aa1e044f8644e0989cc63220", + "method": "Median of 3 round means, 10 launches per round; unchanged msprof harness with decimal 8 GB L2 flush and alternating implementation order." + } +} diff --git a/examples/ascend/flashmla/results/staged-decode-compiler-equivalence-2026-10-18.json b/examples/ascend/flashmla/results/staged-decode-compiler-equivalence-2026-10-18.json new file mode 100644 index 0000000000..4fda938e1d --- /dev/null +++ b/examples/ascend/flashmla/results/staged-decode-compiler-equivalence-2026-10-18.json @@ -0,0 +1,633 @@ +{ + "status": "PASS", + "completed_cases": 23, + "expected_cases": 23, + "source_sha256": "f93885c71bba1103c1b98569a1cad4c5a56d044e7d47207f5b502ca26ed1aa4e", + "native_sha256": "22613b62b8e1c783b8c5592d808fa62ef8466c99613ea34d83a5b751ab832112", + "compiler_commit": "34894d2ee0bf1c311e3587dd530ddb989ae8fa70", + "measurement_compiler_commit": "0d86dd0b886624bb2e35bbfdf83e9e30cc358c98", + "scope": "Host compilation and fresh-process cache-hit verification only. Every complete executable and its host source, device source, and parameter file match the historical device-tested baseline byte for byte.", + "cases": [ + { + "case": "perf-decode-0", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 105364, + "sha256": "3f182b885d5105f0b135f294f3c99a4dfd28404e7332dafabe0c59a3b52ff9af", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "84e6d5599db6daf18ca59232ad82561e728cc61a31cd08ae60e03415d0414457", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60356, + "sha256": "7582cfe0527c3b7ccc01702722b9cce149e32367fb8dce0067dd47c55b64ace9", + "byte_exact": true + }, + "params.json": { + "bytes": 2862, + "sha256": "d5a314ba8db2a69f313bab7347872212d1f9c46e97c21f226e15d4d76e562211", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-1", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 105400, + "sha256": "0d646f68dd381cc577402f82c3f5bc5455728abc9ca31bef7666fa1e6e01e1f4", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "591fa54ff4571f5faad6e6d86a89bbde6b1d216e6ef3e43911c4c01fbac8ca72", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60356, + "sha256": "a3a0dd5de1452b50b5c5a69ee76f256cedbf451e45b5c0b288f3e66beb3e4f83", + "byte_exact": true + }, + "params.json": { + "bytes": 2809, + "sha256": "e515f6c202ccd4cde465c188840f4490791b2ed3ea1a4f33ab7fea218e66e08c", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-2", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 105400, + "sha256": "58ea1413acf70fcc352190b48a0663173a05cb68bffa059a487637ffa7ce0c40", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "c1e6b2a804b993b324c9ff5766a22a83039317a8ab8967fac8c812c332007b4c", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60372, + "sha256": "01289211fb6626b2861cfb49b674f7b18b4ecdca2676300dffc4067864de6c42", + "byte_exact": true + }, + "params.json": { + "bytes": 2864, + "sha256": "dbf841db05620f5711b86b0807713403f1ee5372dbc328657ce74d7a59a8f50e", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-3", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 105400, + "sha256": "20f3b161427a3e863eed3d1a61cc6e68d51c137d815b85c2d1697e847cb2ee13", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "64c8fe00c66cdeffce87e76e6c78a28f0501677aaa7d806ccf066639ef0a083d", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60374, + "sha256": "67f1b1503749629d20878d9a6fcd0995ac3561a2276a9d3f068baa8ba49f4a4b", + "byte_exact": true + }, + "params.json": { + "bytes": 2865, + "sha256": "038917e697b9bcc5efd07aeb93cc48cd4d7da9d0dc8290696ebd74b4c1d18013", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-4", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 111833, + "sha256": "4e55816d2dec25b4a1c9c516cc494f23dcd43369c60b465682102e283d970e2b", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "5ef3da32809002d3aa30233dc1178af0bb08adaf29bc79a2d178b8429cfd6e71", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60354, + "sha256": "2b19a9354c0f4e1bfeac6d340b8866ec2d7795cf538b873f42967efd6a9b34d5", + "byte_exact": true + }, + "params.json": { + "bytes": 2861, + "sha256": "a84e210bbd51b76b4b019d0f6659e6fe0dc5a3e7f159c226652182c016f1bc23", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-5", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 111869, + "sha256": "555a06d258fc7240aac3e1c619b1528c441ae96de766fd261febd131d5fe6016", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "2fed385fe047ee3b00cee4862d562f4f3731ca5069b23a0ef3d9bcdc47a8a18f", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60356, + "sha256": "7681a7cd8d92b8274fd4b907e29c69b283cdc84454b890266f117eddb8de49a0", + "byte_exact": true + }, + "params.json": { + "bytes": 2809, + "sha256": "bc6e53a560ab7684e31c6f5f3c9b4e020efc0a52ff4801dc6e58496466d7b3a6", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-6", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 111869, + "sha256": "d65c768511aca6d4d976fbadad83f3851d6af32c45ab52864844a3ccf4f3c9d0", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "e2a86bcd3ea646947bbb08065cc768a0d6747b24be15924e8cf892176a790770", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60372, + "sha256": "4dd59184755da84b04d09158324fb02684948283e52746ebecea1154895fc04a", + "byte_exact": true + }, + "params.json": { + "bytes": 2864, + "sha256": "d524ba9e5f24a6071ffafdf6708c76d2af50f3570938e92fc13bfac49bac2e90", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-7", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 111869, + "sha256": "39b24eeae366e9e57a5202dfc9270143169fbc3da2eb6d6c7cc70383669289a6", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "fe0d24847c7bf4735006ba7ede767cc7818a24bb98db3ab0ad2db8cd71b598d0", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60372, + "sha256": "74beeec3d688b78e2cbfcee5d113a7fb1ed3c57c719a75aaf8c49dfe0096b364", + "byte_exact": true + }, + "params.json": { + "bytes": 2864, + "sha256": "26dcb19de9551a0592cf84c10dd253110771dc7d23f57ef13a342d026276ed00", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-8", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 105436, + "sha256": "5be6061dce1d74f9b648bd6fbf52031d2dbe7169339851417a98fa14722753a4", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "af6c14d547ede8f7d14e85dc602a02aa05377a51fb58289f0deb399ab3e5c013", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60364, + "sha256": "04832a6a3ad52ed8fa0412894f93c0bf0e2decb87205f1c19c8a63e768f871cf", + "byte_exact": true + }, + "params.json": { + "bytes": 2803, + "sha256": "c01ab086bd7771473683ebb5947a425938a778ff9910a7e02a67a7652f78b203", + "byte_exact": true + } + } + }, + { + "case": "perf-decode-9", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 111905, + "sha256": "6a5d45879fed20f4019e5ffa19d3ad37b479cb85d09599f544d0e4df3069a7ee", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "be18457d3269f3dd4f90d89c23d88cf7384fb0361d98ec4b30b6a6e1b060a29b", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60364, + "sha256": "e2076bdeb21b05920833d8c9ed1d5aa4d9fe049265fed60d82a8cce724cab7bd", + "byte_exact": true + }, + "params.json": { + "bytes": 2860, + "sha256": "f0a5605b652aa67bcd6f730006e0d2a37cffb8bf87804d28b3d9872ea1017f07", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-0", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 97643, + "sha256": "0ad1ea20e17c87a94ffb52c12e6f08a01f37b0e9adb0864aba80a24b4ef083b6", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "3081a2a724d690b948464d37874d3252fbe16940b5fc3e2da8b5c7cecd4c2b2c", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 58553, + "sha256": "9c4f71285635719cdd36dbb259496627455e696a67d009308a42d1069da8687c", + "byte_exact": true + }, + "params.json": { + "bytes": 2852, + "sha256": "94729d9223242e5c03962eadf44ebc94979a14b2ddcddbdc18af13975cb92beb", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-1", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 101442, + "sha256": "92894373cc2a837955af5432294f9d152c44dd9c669bddb228573cde737136e7", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "8d0fb2103260a53cb85d7c362eaf1423ecd9d52fff00993c5d23c02ff365bdc6", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60336, + "sha256": "a00ea8a9f7539186f05b656ad0a803e18bb985aebd41ca7d660a184fba18bb7d", + "byte_exact": true + }, + "params.json": { + "bytes": 2859, + "sha256": "1f47371b2a141d3625d94c5e4e87b5cf7b4d5f6493576e9faf2be5ca78399595", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-2", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 107786, + "sha256": "d814f4a9dd22ee06d601060941fd692e925900619e9792c4bd163a7167d5f516", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "5c24323e386d36ba414b113959e4c1e82b9cca152ed9e9750868bb478c9675aa", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60336, + "sha256": "b0958f7e793451cad9b2f1bfa1d6d7bce5b37da8d3790ccf99e8bd3b10cc319f", + "byte_exact": true + }, + "params.json": { + "bytes": 2859, + "sha256": "746ff6c0e045bd0f25f8208aca39f37eb91766c821bd5de346ddeeacd687102d", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-3", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 101329, + "sha256": "4af566bd786c638ad462ae9e138f26bab28c3c914fe32922f97d073c5b734feb", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "6eb10f40cefe97202945f9833fb6250b8e34bfc86a9134f7f9f0b86118efa62a", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60009, + "sha256": "cb0dd345f72f51cf001f91e0654086909de260dce63e19a023624394201b3d3f", + "byte_exact": true + }, + "params.json": { + "bytes": 2855, + "sha256": "aa9c3d468b4c1209d6337f643f10fd390abfce90cca32212d40a2dbaf3a7b54a", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-4", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 109589, + "sha256": "bb4189d1d1dde66ee6dfecdc4d91439d496754284982984b760cba581c1e27bf", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "3671b825e0a53435f50329ff5eff5ffc5ee1df819df7967ee9cde560fa998f37", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60897, + "sha256": "8da2d08d3096cd5df3794df8a4f68fcfb98e97cab34c5b2d0a992bf0b7a37963", + "byte_exact": true + }, + "params.json": { + "bytes": 2861, + "sha256": "55b0ecd4268cced9151f70f8d00399f5a23d013810cf119e5016618aa774bb31", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-5", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 107790, + "sha256": "5edfa1f0c7db7538623dc09e7fb86f040b293a9e528fd482203e993ecfcfacc8", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "5d6dffe740fa0c28cd96b2782d3e9242eca35b6943cb0129ee27dbf54b344ed3", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60044, + "sha256": "c070f8da4186926f6fd570420e52e821eb107b64d981cb1611a0d27dc5af554c", + "byte_exact": true + }, + "params.json": { + "bytes": 2861, + "sha256": "06e5d5e49db8db55f1beb8d43c023efffe4a195febb07ced6ff8f740f023fb3b", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-6", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 107612, + "sha256": "f09c7e08504d16ac21f1f506317e4f843496032322efdfe7799e403b71f79d11", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "c347793c8a4d99a65cfad592757bdb445b250fffffec6aefcdaaf2896a26f307", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60007, + "sha256": "4f4874ff1ba3e93e2d49582b8e5ba917a079291963123dfbda1b186d74769140", + "byte_exact": true + }, + "params.json": { + "bytes": 2854, + "sha256": "46a61247a5e9390efa77afee7c96334b274acbad63be7a6d16a8d67c4c4b1ab9", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-7", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 115867, + "sha256": "69a02cba0945aef1a9358a311f4145f996037295bdbc0f88d7a947f1a8fb2d1a", + "byte_exact": true + }, + "executable.so": { + "bytes": 135880, + "sha256": "f27002b699edc27ac145a3380e4d867b7df2972ca055a4d28122619200dab9aa", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60897, + "sha256": "167912293d2f0a6bec374adaac0fe9dbe3c22d09567efa38902daa5ce0662df7", + "byte_exact": true + }, + "params.json": { + "bytes": 2861, + "sha256": "62db1a54698793add527e0843a8703a51f35addd5f6ddcac984ac3ba5a25fe5b", + "byte_exact": true + } + } + }, + { + "case": "correct-decode-8", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 114167, + "sha256": "aa6fe6f6c65d17f2a4af8b6b7ee265045f974462c165d0fb1026d7a7ffe875d4", + "byte_exact": true + }, + "executable.so": { + "bytes": 135880, + "sha256": "9c38f3417e390bc81f481364c2e300f361c36153478895f6f8e5d8fa552927a8", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60042, + "sha256": "0ed5752689d6bb4b6ed00c01546c9253bea34434fdea5bf104d0eba5e63d3be7", + "byte_exact": true + }, + "params.json": { + "bytes": 2860, + "sha256": "746380a3d7bb6458fc0f70bbde299f51c59bdc8f95854666a86788394e3c8041", + "byte_exact": true + } + } + }, + { + "case": "boundary-decode-0", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 107075, + "sha256": "dd6fb499a6654bcaaba3f272a784576bfaaf99f6a0ea882d1d07769de8f88050", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "859c08dc36d4f7b98b5d52e08b4237571db63bfdf90f15cffa5cbef21d45b4ae", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60645, + "sha256": "2c8799d1b3443e07f07bb62be08144694de7bc65d60568f1adb1f8aecbd6e64f", + "byte_exact": true + }, + "params.json": { + "bytes": 2914, + "sha256": "9cb5ecb4d1a7baf7cdfa9b0c2f524f151879d3912f1ded4543b4ad0fef1eb953", + "byte_exact": true + } + } + }, + { + "case": "boundary-decode-1", + "status": "PASS", + "baseline": "v91", + "files": { + "device_kernel.asc": { + "bytes": 107096, + "sha256": "7d230e386970e45785eed8336fe1399fb22cfb437a41bdb8e949e09408510212", + "byte_exact": true + }, + "executable.so": { + "bytes": 70344, + "sha256": "40afcbee81d771f2276a16476bfeaa5772d2d39313f7cb5e97293794ee8fcb77", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60610, + "sha256": "11b0de30e2b6ecd06b67671f4dd1a675c45423b9142dcf552d14b9431e10a871", + "byte_exact": true + }, + "params.json": { + "bytes": 2914, + "sha256": "9cb5ecb4d1a7baf7cdfa9b0c2f524f151879d3912f1ded4543b4ad0fef1eb953", + "byte_exact": true + } + } + }, + { + "case": "boundary-decode-2", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 113091, + "sha256": "55f07ca2e19041ad53fc055a0f5a1f39f5514ecba6f6f2e1f45e8d40fabff67f", + "byte_exact": true + }, + "executable.so": { + "bytes": 135880, + "sha256": "2cae8754ea7a6e8c8a2bcc3722cf7fc1d6a551bbf36f2101c29808b9fd61a21d", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60645, + "sha256": "f4a49fe73206259166424470e9253687a9c733bc2dbf96cb5629058c0aff75d2", + "byte_exact": true + }, + "params.json": { + "bytes": 2914, + "sha256": "4765edbf2a3716f1a1c859198fccb8480be80dfefcf0252485756a9f8c6404ae", + "byte_exact": true + } + } + }, + { + "case": "boundary-decode-3", + "status": "PASS", + "baseline": "v100", + "files": { + "device_kernel.asc": { + "bytes": 113112, + "sha256": "86ab9683466f7c9470772d7b6d3870f90b69a0b4ecf06a574bd6923310765b12", + "byte_exact": true + }, + "executable.so": { + "bytes": 135880, + "sha256": "d0056be9e191fb2db0125d6fda50fb5b5cc9673b44ef03acacdc986d954f9032", + "byte_exact": true + }, + "host_kernel.c": { + "bytes": 60610, + "sha256": "cb5c35ce0c5ee2334795c654c5645c6c24af3804dc1374896972a51c4b2ebee3", + "byte_exact": true + }, + "params.json": { + "bytes": 2914, + "sha256": "4765edbf2a3716f1a1c859198fccb8480be80dfefcf0252485756a9f8c6404ae", + "byte_exact": true + } + } + } + ] +} diff --git a/examples/ascend/flashmla/results/staged-main-rebase-2026-10-20.json b/examples/ascend/flashmla/results/staged-main-rebase-2026-10-20.json new file mode 100644 index 0000000000..bc281e1670 --- /dev/null +++ b/examples/ascend/flashmla/results/staged-main-rebase-2026-10-20.json @@ -0,0 +1,275 @@ +{ + "date": "2026-10-20", + "recorded_at_utc": "2026-10-20T10:26:51.598797+00:00", + "scope": "Local main rebase, incremental native build, host tests, and source-only FlashMLA lowering; no NPU correctness or performance rerun.", + "worktree": "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example", + "branch": "codex/flashmla-manual-phase-perf", + "previous_head": "9c46d1cdab152218a9335fadfcde8816409f402b", + "upstream_main": "876a59e6eb8091ea608c14a84d5300190c3735ce", + "head": "00e7b9bb64c892c5323f42bede764e8940ceb4cd", + "backup_branch": "flashmla-before-main-rebase-20261020", + "original_merge_base": "adeea32bcecdddca32195a95e90730b175606a0e", + "local_commit_count": 22, + "rewritten_commits": [ + { + "old": "a638c2ac356e6dbdf22ccb5b8889ef2ffc8eb32b", + "new": "a864bfdb0e07d6c54476f5d20f64b2ec92d9ddab" + }, + { + "old": "0a9ec706796aefb4b601395646e2f673f76df904", + "new": "9d0090bf3956cb3cf57323a6e249080cd30669ad" + }, + { + "old": "1ba2325dc45e7484b5b655c121e159c60905377d", + "new": "ca11cac23165135fc206d0c338872e9544e69efe" + }, + { + "old": "64e72d5a41129d5d632f459edd07144ea19a19b7", + "new": "bd25ad6c12bbbdcb6191f74bd28db1f7a7f03dab" + }, + { + "old": "e4799b664388034b49c42d027e84b1567e39d6a9", + "new": "c114cb8f11e33e8cfebd79a6ede7abac0f35cf91" + }, + { + "old": "b683a26a8f13ffa527f33e751473c400be70c379", + "new": "2d3e1f8603243e87c533e0c7296d8c519b60732c" + }, + { + "old": "2dbae87c4daa435052ca19c2490a05a911f48404", + "new": "0868fb0a110f7ed9ae97c9a7bd392a5360bd554d" + }, + { + "old": "29e8254b689dfb3c5d9a178b0920af257baca96b", + "new": "706160330cb9ba51c8af7e47bececc309badf042" + }, + { + "old": "c1b854b39c5b43a7b1300d9c97770763620dd2b9", + "new": "2ee1a2f42ef8618d3d0f3022e9991e62046d8d68" + }, + { + "old": "caa62b67226fe1bcac715bad97bfadd0cdc10113", + "new": "fcbc968f83adfddc8c8de440d8e191d145d0018c" + }, + { + "old": "35a66f40f309d28a34dd41c9155381488086a97a", + "new": "d3042c789b64111d327f32381a246dd38c15fe83" + }, + { + "old": "26799f47e03a3b9e5dddd99039ac11ab8ed325ad", + "new": "dbb5728bcee850134bb1cc53038e6330544c112f" + }, + { + "old": "9252772bb93ff83dc930bd666daefa22e86fd497", + "new": "d6106572a566105d5117444096e103426b7d3de9" + }, + { + "old": "2f2e768d29d352c7c296db03164c9d7032ba9a6f", + "new": "feb8e9f71a76e10cd0631df89e57ee9822e4d66c" + }, + { + "old": "cbbea6969ab4c873e46476d03ac1274f8ac0ff55", + "new": "d687d8cee7fdc81a8559bf4783b1bfe18a0025c2" + }, + { + "old": "0d86dd0b886624bb2e35bbfdf83e9e30cc358c98", + "new": "20929efc503511f51770c594179371f56a5da5b5" + }, + { + "old": "34894d2ee0bf1c311e3587dd530ddb989ae8fa70", + "new": "ac38d6568dc367809fedbf8076af620ac06f6d49" + }, + { + "old": "218c9c9131511952aa7b8e7e2b156e5c4bf2620c", + "new": "9e65a4d97054b414bc1bdd4209f88d2507cd6e0f" + }, + { + "old": "c536475fa3db8e41e04546388b6c76e15681190c", + "new": "c206829bcf06260ea7ec206122fb16eb268962dd" + }, + { + "old": "9e691b088b49567afe0d4e4690a5e94983e75f84", + "new": "48e8d3aa9c8c35de4fa1f7623e13ddc33c634cf3" + }, + { + "old": "e9244819b52faa5775dc777570ff9110d0234ba1", + "new": "80d6fb91facd1281dfc467eb2ac6312739165416" + }, + { + "old": "9c46d1cdab152218a9335fadfcde8816409f402b", + "new": "00e7b9bb64c892c5323f42bede764e8940ceb4cd" + } + ], + "upstream_features_removed_from_local_delta": { + "7617c7d": "SS buffer and automatic versions", + "390ecf3": "Independent AIV execution", + "0bd33f9": "Partial L0 copies and explicit pitch", + "f16f1d7": "SIMD fill in multibuffered regions", + "51bdf54": "SIMD reduction row stores", + "c55eeb1": "PTO execution and codegen routing" + }, + "remaining_prerequisites": { + "old_commit": "35a66f40f309d28a34dd41c9155381488086a97a", + "commit": "d3042c789b64111d327f32381a246dd38c15fe83", + "old_changed_files": 77, + "new_changed_files": 22, + "scope": "NZ post-copy pitches, compact allocation, physical access footprints, and views. The composite integration commit could not be dropped as a whole." + }, + "preservation": { + "committed_flashmla_directory_unchanged": true, + "uncommitted_documentation_restored_before_further_updates": { + "docs/programming_guides/ascend_issue_phase.md": "993e2b7674a5e455b45c48f91cc8d2c087f10bb97c5fd6665e8b1e5d70da5268", + "examples/ascend/flashmla/README.md": "0cc37ecbc6c2a20a16c21c6c793ee56e87760c49fdb7bd9200d2d14f15c2f6b9", + "examples/ascend/flashmla/RESULTS.md": "19d25f83873d91ad6a9baa145004ae2502fb9b1e2c72ef30cc14444788cbc46c", + "examples/ascend/flashmla/WORKTREES.md": "bdb979c46fb26688144bb9722abaccb0c270a3dcca8b5921b1ea59ed3b991a77", + "examples/ascend/flashmla/results/staged-manual-order-2026-10-20.json": "fb47770a7bcc9c7e7ea5a643c786785366b3698377a9134416587ba1223f05bb" + }, + "documentation_stash": "b90ec439cd078280f38609dadf2c84272b0626a0", + "recursive_dependencies": [ + { + "path": "3rdparty/cutlass", + "expected": "b2dd65dc864e09688245b316ac46c4a6cd07e15c", + "head": "b2dd65dc864e09688245b316ac46c4a6cd07e15c", + "status": "" + }, + { + "path": "3rdparty/tvm", + "expected": "4211874e9e4b8770fd20ad0d98ab80afd1df5028", + "head": "4211874e9e4b8770fd20ad0d98ab80afd1df5028", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/OpenCL-Headers", + "expected": "b590a6bfe034ea3a418b7b523e3490956bcb367a", + "head": "b590a6bfe034ea3a418b7b523e3490956bcb367a", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/cutlass", + "expected": "b2dd65dc864e09688245b316ac46c4a6cd07e15c", + "head": "b2dd65dc864e09688245b316ac46c4a6cd07e15c", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/cutlass_fpA_intB_gemm", + "expected": "72b9883c986a2ff427ca61ac0b14ad59be1dc862", + "head": "72b9883c986a2ff427ca61ac0b14ad59be1dc862", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/cutlass_fpA_intB_gemm/cutlass", + "expected": "cc85b64cf676c45f98a17e3a47c0aafcf817f088", + "head": "cc85b64cf676c45f98a17e3a47c0aafcf817f088", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/libflash_attn", + "expected": "07ba35bae96900e51b4d63ef3487ba24850da870", + "head": "07ba35bae96900e51b4d63ef3487ba24850da870", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/libflash_attn/cutlass", + "expected": "e0aaa3c3b38db9a89c31f04fef91e92123ad5e2e", + "head": "e0aaa3c3b38db9a89c31f04fef91e92123ad5e2e", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/tvm-ffi", + "expected": "3c35034fd1026011736e19a4e0e1ed0f22058c42", + "head": "3c35034fd1026011736e19a4e0e1ed0f22058c42", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/dlpack", + "expected": "84d107bf416c6bab9ae68ad285876600d230490d", + "head": "84d107bf416c6bab9ae68ad285876600d230490d", + "status": "" + }, + { + "path": "3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/libbacktrace", + "expected": "793921876c981ce49759114d7bb89bb89b2d3a2d", + "head": "793921876c981ce49759114d7bb89bb89b2d3a2d", + "status": "" + } + ] + }, + "native_build": { + "command": "source .tlcli/env.sh && cmake --build build --target tilelang --parallel 8", + "result": "PASS", + "incremental": true, + "libraries": { + "build/lib/libtilelang.so": "489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b", + "build/lib/libtvm_compiler.so": "572a2285726e0d0565f47d4e56d757d31c956959605ff126d4725f8a510004a4", + "build/lib/libtvm_runtime.so": "2b6cd5655295f8cd5915abcb39929aa4f0a2522579d2a6440c5503820e958b00" + }, + "log": ".tlcli/rebase-main-20261020/build.log" + }, + "host_tests": { + "command": "python -m pytest testing/ascend/auto_schedule testing/ascend/transform/test_ascend_lower_copy.py testing/ascend/transform/test_ascend_insert_nd2nz.py testing/ascend/transform/test_ascend_l1_offsets.py testing/ascend/layout/test_ascend_fractal_storage_integration.py::test_nz_fixpipe_layout_survives_full_lowering testing/ascend/layout/test_ascend_fractal_storage_integration.py::test_nz_fixpipe_does_not_apply_nd_tail_padding testing/ascend/layout/test_ascend_fractal_storage_integration.py::test_padded_alias_storage_sets_the_physical_version_pitch testing/ascend/layout/test_ascend_fractal_storage_integration.py::test_dynamic_fixpipe_padding_survives_full_lowering testing/ascend/layout/test_tilelang_ascend_padded_nz.py::test_compact_nz_layout_keeps_stages_tightly_packed testing/ascend/language/test_tilelang_ascend_pto_mad_direction.py -m 'not ascend' -n 4 -q --maxfail=1", + "environment": { + "setup": "source .tlcli/env.sh", + "ASCEND_RT_VISIBLE_DEVICES": "", + "OMP_NUM_THREADS": "1", + "TILELANG_DISABLE_CACHE": "1" + }, + "result": "PASS", + "passed": 987, + "warnings": 114, + "seconds": 26.33, + "log": ".tlcli/rebase-main-20261020/host-tests-final.log", + "initial_selection_failure": { + "passed_before_stop": 970, + "failed": 3, + "cause": "Unmarked NPU tests in the two layout files bypassed the not-ascend marker filter and failed at aclInit (107001, Invalid device) with NPU visibility disabled. The final run explicitly selects only host tests from these files.", + "log": ".tlcli/rebase-main-20261020/host-tests.log" + } + }, + "kernel_lowering": [ + { + "case": "bf16", + "nq": 65, + "nk": 2048, + "topk": 640, + "kwargs": { + "variable_lengths": true + }, + "tilelang": "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/tilelang/__init__.py", + "source_sha256": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2", + "native_library_sha256": "489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b", + "validation": "compile-only: full TileLang lowering and ASC source generation; no vendor compilation or NPU execution", + "result": "PASS", + "seconds": 298.5134978940041, + "generated_source_sha256": "1d5e14584e51bd0212411e17d48b786f2a582bdd9c5811a5e071b15aa63450f6" + }, + { + "case": "fp4", + "nq": 65, + "nk": 2048, + "topk": 128, + "kwargs": { + "variable_lengths": true, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "variable_extra_lengths": true, + "kv_storage_bytes": 1081344, + "extra_storage_bytes": 589824 + }, + "tilelang": "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/tilelang/__init__.py", + "source_sha256": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2", + "native_library_sha256": "489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b", + "validation": "compile-only: full TileLang lowering and ASC source generation; no vendor compilation or NPU execution", + "result": "PASS", + "seconds": 345.78024685400305, + "generated_source_sha256": "51a7c57bb817d7f9d8eb8168d73bc8f44a414cf255096d4c2e5cf4f5444fe7eb" + } + ], + "kernel_lowering_warnings": "Both cases reported intermediate intra-core flag-id high-water marks of 11 for MTE1_MTE2 and MTE2_MTE1; the later RewriteFlagToBuf pass completed and final ASC generation succeeded. This is not device execution evidence.", + "limitations": [ + "No Bisheng compilation of the generated FlashMLA kernels in this validation.", + "No NPU execution, correctness matrix, or performance matrix rerun on the new compiler.", + "Earlier device and timing evidence retains its original compiler/library identity even though the kernel source is unchanged." + ], + "artifacts": ".tlcli/rebase-main-20261020/" +} diff --git a/examples/ascend/flashmla/results/staged-main-rebase-decode-bench-2026-10-20.json b/examples/ascend/flashmla/results/staged-main-rebase-decode-bench-2026-10-20.json new file mode 100644 index 0000000000..28a3501e0f --- /dev/null +++ b/examples/ascend/flashmla/results/staged-main-rebase-decode-bench-2026-10-20.json @@ -0,0 +1,3079 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "6", + "cases": [ + { + "id": 0, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 78.63044738769531, + 78.77349853515625, + 78.84502410888672 + ], + "launch_us": [ + [ + 78.91654968261719, + 78.91654968261719, + 79.87022399902344, + 78.91654968261719, + 77.96287536621094, + 77.72445678710938, + 79.63180541992188, + 77.48603820800781, + 78.2012939453125, + 78.67813110351562 + ], + [ + 79.15496826171875, + 78.91654968261719, + 79.39338684082031, + 79.15496826171875, + 78.43971252441406, + 78.43971252441406, + 78.67813110351562, + 78.43971252441406, + 78.43971252441406, + 78.67813110351562 + ], + [ + 79.87022399902344, + 79.39338684082031, + 78.43971252441406, + 77.96287536621094, + 78.2012939453125, + 78.43971252441406, + 79.63180541992188, + 78.43971252441406, + 79.39338684082031, + 78.67813110351562 + ] + ], + "median_us": 78.77349853515625, + "tflops": 272.6149895502722 + }, + "manual": { + "round_mean_us": [ + 77.12841033935547, + 77.056884765625, + 77.36682891845703 + ], + "launch_us": [ + [ + 76.77078247070312, + 75.57868957519531, + 77.24761962890625, + 76.77078247070312, + 77.48603820800781, + 76.77078247070312, + 78.91654968261719, + 77.48603820800781, + 78.43971252441406, + 75.81710815429688 + ], + [ + 78.67813110351562, + 77.48603820800781, + 77.48603820800781, + 76.53236389160156, + 77.00920104980469, + 75.57868957519531, + 76.53236389160156, + 77.48603820800781, + 76.77078247070312, + 77.00920104980469 + ], + [ + 77.48603820800781, + 78.43971252441406, + 76.53236389160156, + 78.2012939453125, + 77.00920104980469, + 76.77078247070312, + 76.77078247070312, + 77.48603820800781, + 78.2012939453125, + 76.77078247070312 + ] + ], + "median_us": 77.12841033935547, + "tflops": 278.42965238766595 + }, + "staged": { + "round_mean_us": [ + 73.36139678955078, + 72.83687591552734, + 72.74150848388672 + ], + "launch_us": [ + [ + 74.14817810058594, + 73.19450378417969, + 72.71766662597656, + 74.86343383789062, + 72.95608520507812, + 74.3865966796875, + 72.71766662597656, + 71.76399230957031, + 72.95608520507812, + 73.90975952148438 + ], + [ + 72.71766662597656, + 72.71766662597656, + 71.28715515136719, + 74.14817810058594, + 72.71766662597656, + 72.71766662597656, + 72.95608520507812, + 73.43292236328125, + 72.24082946777344, + 73.43292236328125 + ], + [ + 73.19450378417969, + 73.19450378417969, + 72.24082946777344, + 72.71766662597656, + 72.95608520507812, + 72.00241088867188, + 72.479248046875, + 72.24082946777344, + 72.71766662597656, + 73.67134094238281 + ] + ], + "median_us": 72.83687591552734, + "tflops": 294.8346728229457 + } + } + }, + { + "id": 1, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 145.48301696777344, + 145.31612396240234, + 145.2922821044922 + ], + "launch_us": [ + [ + 144.95849609375, + 146.1505889892578, + 144.48165893554688, + 146.38900756835938, + 146.8658447265625, + 146.62742614746094, + 144.72007751464844, + 145.91217041015625, + 143.52798461914062, + 145.19691467285156 + ], + [ + 145.43533325195312, + 144.2432403564453, + 143.52798461914062, + 146.38900756835938, + 143.52798461914062, + 145.6737518310547, + 147.34268188476562, + 146.62742614746094, + 145.19691467285156, + 145.19691467285156 + ], + [ + 144.00482177734375, + 144.48165893554688, + 144.48165893554688, + 146.8658447265625, + 145.43533325195312, + 144.72007751464844, + 145.6737518310547, + 146.1505889892578, + 146.38900756835938, + 144.72007751464844 + ] + ], + "median_us": 145.31612396240234, + "tflops": 295.56027086926963 + }, + "manual": { + "round_mean_us": [ + 143.0988311767578, + 143.31340789794922, + 142.7173614501953 + ], + "launch_us": [ + [ + 141.62063598632812, + 142.09747314453125, + 143.52798461914062, + 143.52798461914062, + 144.72007751464844, + 143.0511474609375, + 144.2432403564453, + 141.62063598632812, + 142.09747314453125, + 144.48165893554688 + ], + [ + 144.00482177734375, + 143.0511474609375, + 143.52798461914062, + 144.2432403564453, + 141.8590545654297, + 143.7664031982422, + 142.09747314453125, + 144.00482177734375, + 142.81272888183594, + 143.7664031982422 + ], + [ + 143.0511474609375, + 142.81272888183594, + 142.3358917236328, + 144.00482177734375, + 141.62063598632812, + 142.57431030273438, + 142.09747314453125, + 141.8590545654297, + 143.52798461914062, + 143.28956604003906 + ] + ], + "median_us": 143.0988311767578, + "tflops": 300.13992851519464 + }, + "staged": { + "round_mean_us": [ + 136.80458068847656, + 137.1622085571289, + 137.3291015625 + ], + "launch_us": [ + [ + 136.6138458251953, + 137.56752014160156, + 136.6138458251953, + 137.80593872070312, + 138.0443572998047, + 136.1370086669922, + 135.18333435058594, + 136.85226440429688, + 136.37542724609375, + 136.85226440429688 + ], + [ + 137.56752014160156, + 136.6138458251953, + 136.85226440429688, + 137.3291015625, + 136.1370086669922, + 137.09068298339844, + 137.09068298339844, + 138.0443572998047, + 138.28277587890625, + 136.6138458251953 + ], + [ + 137.09068298339844, + 137.09068298339844, + 137.3291015625, + 137.09068298339844, + 137.80593872070312, + 137.80593872070312, + 137.56752014160156, + 138.5211944580078, + 135.4217529296875, + 137.56752014160156 + ] + ], + "median_us": 137.1622085571289, + "tflops": 313.1305146789846 + } + } + }, + { + "id": 2, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 275.5880355834961, + 275.9695053100586, + 275.8502960205078 + ], + "launch_us": [ + [ + 275.3734588623047, + 274.89662170410156, + 275.3734588623047, + 276.0887145996094, + 275.1350402832031, + 276.80397033691406, + 274.658203125, + 275.8502960205078, + 276.32713317871094, + 275.3734588623047 + ], + [ + 274.658203125, + 277.0423889160156, + 277.51922607421875, + 276.5655517578125, + 274.89662170410156, + 274.658203125, + 277.0423889160156, + 277.2808074951172, + 275.1350402832031, + 274.89662170410156 + ], + [ + 277.0423889160156, + 274.89662170410156, + 275.61187744140625, + 275.1350402832031, + 275.61187744140625, + 277.0423889160156, + 276.32713317871094, + 275.8502960205078, + 275.61187744140625, + 275.3734588623047 + ] + ], + "median_us": 275.8502960205078, + "tflops": 311.39841848715616 + }, + "manual": { + "round_mean_us": [ + 271.3441848754883, + 271.7256546020508, + 271.63028717041016 + ], + "launch_us": [ + [ + 271.3203430175781, + 270.3666687011719, + 271.5587615966797, + 271.5587615966797, + 271.5587615966797, + 272.2740173339844, + 271.5587615966797, + 271.08192443847656, + 270.60508728027344, + 271.5587615966797 + ], + [ + 272.2740173339844, + 272.0355987548828, + 270.60508728027344, + 271.3203430175781, + 271.3203430175781, + 270.843505859375, + 271.79718017578125, + 272.0355987548828, + 270.843505859375, + 274.1813659667969 + ], + [ + 272.0355987548828, + 271.08192443847656, + 272.0355987548828, + 270.843505859375, + 271.08192443847656, + 271.3203430175781, + 272.7508544921875, + 271.5587615966797, + 272.2740173339844, + 271.3203430175781 + ] + ], + "median_us": 271.63028717041016, + "tflops": 316.2362592729217 + }, + "staged": { + "round_mean_us": [ + 264.5254135131836, + 264.4538879394531, + 264.66846466064453 + ], + "launch_us": [ + [ + 261.54518127441406, + 264.88304138183594, + 263.45252990722656, + 265.35987854003906, + 266.5519714355469, + 262.0220184326172, + 263.9293670654297, + 264.16778564453125, + 269.88983154296875, + 263.45252990722656 + ], + [ + 262.97569274902344, + 262.7372741699219, + 263.6909484863281, + 264.6446228027344, + 263.9293670654297, + 264.16778564453125, + 265.1214599609375, + 265.1214599609375, + 266.5519714355469, + 265.5982971191406 + ], + [ + 264.6446228027344, + 263.9293670654297, + 266.07513427734375, + 264.16778564453125, + 264.4062042236328, + 263.9293670654297, + 266.3135528564453, + 265.35987854003906, + 262.0220184326172, + 265.8367156982422 + ] + ], + "median_us": 264.5254135131836, + "tflops": 324.7300317166648 + } + } + }, + { + "id": 3, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 541.4009094238281, + 541.9492721557617, + 541.0909652709961 + ], + "launch_us": [ + [ + 544.0711975097656, + 544.5480346679688, + 541.4485931396484, + 542.6406860351562, + 540.9717559814453, + 540.2565002441406, + 539.0644073486328, + 540.0180816650391, + 540.0180816650391, + 540.9717559814453 + ], + [ + 544.7864532470703, + 541.68701171875, + 541.4485931396484, + 541.2101745605469, + 540.7333374023438, + 540.7333374023438, + 540.9717559814453, + 544.0711975097656, + 542.8791046142578, + 540.9717559814453 + ], + [ + 540.7333374023438, + 538.1107330322266, + 540.7333374023438, + 540.0180816650391, + 541.4485931396484, + 538.3491516113281, + 545.0248718261719, + 544.3096160888672, + 541.4485931396484, + 540.7333374023438 + ] + ], + "median_us": 541.4009094238281, + "tflops": 317.3225032496386 + }, + "manual": { + "round_mean_us": [ + 534.9159240722656, + 535.0828170776367, + 535.0351333618164 + ], + "launch_us": [ + [ + 533.3423614501953, + 537.872314453125, + 535.7265472412109, + 534.0576171875, + 533.3423614501953, + 535.4881286621094, + 535.4881286621094, + 533.5807800292969, + 534.0576171875, + 536.2033843994141 + ], + [ + 534.5344543457031, + 537.6338958740234, + 534.7728729248047, + 534.2960357666016, + 534.7728729248047, + 535.4881286621094, + 535.9649658203125, + 534.7728729248047, + 535.2497100830078, + 533.3423614501953 + ], + [ + 534.2960357666016, + 536.4418029785156, + 535.7265472412109, + 534.0576171875, + 537.872314453125, + 533.3423614501953, + 535.0112915039062, + 533.5807800292969, + 535.2497100830078, + 534.7728729248047 + ] + ], + "median_us": 535.0351333618164, + "tflops": 321.09796371787326 + }, + "staged": { + "round_mean_us": [ + 518.9418792724609, + 519.5379257202148, + 520.014762878418 + ], + "launch_us": [ + [ + 521.1830139160156, + 516.4146423339844, + 518.3219909667969, + 519.5140838623047, + 520.4677581787109, + 518.798828125, + 519.2756652832031, + 519.2756652832031, + 518.0835723876953, + 518.0835723876953 + ], + [ + 521.1830139160156, + 521.4214324951172, + 517.3683166503906, + 520.2293395996094, + 523.5671997070312, + 518.0835723876953, + 519.2756652832031, + 518.3219909667969, + 519.0372467041016, + 516.8914794921875 + ], + [ + 520.4677581787109, + 518.798828125, + 519.9909210205078, + 519.7525024414062, + 518.5604095458984, + 519.0372467041016, + 518.5604095458984, + 520.2293395996094, + 522.3751068115234, + 522.3751068115234 + ] + ], + "median_us": 519.5379257202148, + "tflops": 330.6759397821483 + } + } + }, + { + "id": 4, + "b": 64, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 70.61958312988281, + 70.64342498779297, + 70.09506225585938 + ], + "launch_us": [ + [ + 69.85664367675781, + 69.85664367675781, + 68.66455078125, + 70.09506225585938, + 70.33348083496094, + 72.00241088867188, + 72.24082946777344, + 69.14138793945312, + 71.28715515136719, + 72.71766662597656 + ], + [ + 73.19450378417969, + 72.24082946777344, + 69.61822509765625, + 70.33348083496094, + 69.37980651855469, + 69.14138793945312, + 72.24082946777344, + 69.14138793945312, + 71.28715515136719, + 69.85664367675781 + ], + [ + 69.37980651855469, + 69.14138793945312, + 69.61822509765625, + 69.14138793945312, + 71.52557373046875, + 71.76399230957031, + 71.28715515136719, + 69.85664367675781, + 69.37980651855469, + 69.85664367675781 + ] + ], + "median_us": 70.61958312988281, + "tflops": 304.09180468403076 + }, + "manual": { + "round_mean_us": [ + 70.54805755615234, + 69.52285766601562, + 70.26195526123047 + ], + "launch_us": [ + [ + 72.24082946777344, + 70.81031799316406, + 69.85664367675781, + 72.24082946777344, + 70.33348083496094, + 68.90296936035156, + 71.52557373046875, + 69.61822509765625, + 68.90296936035156, + 71.04873657226562 + ], + [ + 69.37980651855469, + 69.14138793945312, + 70.81031799316406, + 71.04873657226562, + 68.66455078125, + 69.85664367675781, + 69.85664367675781, + 68.90296936035156, + 68.66455078125, + 68.90296936035156 + ], + [ + 69.61822509765625, + 70.09506225585938, + 69.14138793945312, + 69.85664367675781, + 70.81031799316406, + 71.28715515136719, + 69.14138793945312, + 71.28715515136719, + 72.00241088867188, + 69.37980651855469 + ] + ], + "median_us": 70.26195526123047, + "tflops": 305.6396082368847 + }, + "staged": { + "round_mean_us": [ + 68.80760192871094, + 67.49629974365234, + 67.92545318603516 + ], + "launch_us": [ + [ + 67.47245788574219, + 68.90296936035156, + 68.18771362304688, + 70.09506225585938, + 69.61822509765625, + 69.61822509765625, + 68.90296936035156, + 67.94929504394531, + 67.94929504394531, + 69.37980651855469 + ], + [ + 66.51878356933594, + 66.7572021484375, + 67.71087646484375, + 67.23403930664062, + 66.51878356933594, + 68.66455078125, + 66.99562072753906, + 67.71087646484375, + 67.47245788574219, + 69.37980651855469 + ], + [ + 67.71087646484375, + 67.47245788574219, + 67.47245788574219, + 66.99562072753906, + 68.66455078125, + 69.14138793945312, + 68.90296936035156, + 68.18771362304688, + 67.23403930664062, + 67.47245788574219 + ] + ], + "median_us": 67.92545318603516, + "tflops": 316.1530099944188 + } + } + }, + { + "id": 5, + "b": 128, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 129.77123260498047, + 129.6520233154297, + 129.34207916259766 + ], + "launch_us": [ + [ + 129.69970703125, + 131.6070556640625, + 130.4149627685547, + 128.74603271484375, + 128.5076141357422, + 130.65338134765625, + 129.69970703125, + 127.31552124023438, + 132.08389282226562, + 128.9844512939453 + ], + [ + 129.46128845214844, + 128.5076141357422, + 129.69970703125, + 129.69970703125, + 129.69970703125, + 129.46128845214844, + 131.84547424316406, + 128.03077697753906, + 130.8917999267578, + 129.22286987304688 + ], + [ + 131.36863708496094, + 130.8917999267578, + 128.5076141357422, + 128.9844512939453, + 127.55393981933594, + 128.9844512939453, + 130.4149627685547, + 128.26919555664062, + 130.65338134765625, + 127.7923583984375 + ] + ], + "median_us": 129.6520233154297, + "tflops": 331.26882143218063 + }, + "manual": { + "round_mean_us": [ + 129.72354888916016, + 130.39112091064453, + 130.1288604736328 + ], + "launch_us": [ + [ + 128.74603271484375, + 129.93812561035156, + 130.8917999267578, + 129.46128845214844, + 129.93812561035156, + 128.9844512939453, + 128.5076141357422, + 131.36863708496094, + 130.8917999267578, + 128.5076141357422 + ], + [ + 129.69970703125, + 129.46128845214844, + 131.84547424316406, + 129.22286987304688, + 129.69970703125, + 130.4149627685547, + 131.13021850585938, + 130.8917999267578, + 131.36863708496094, + 130.17654418945312 + ], + [ + 136.85226440429688, + 131.13021850585938, + 128.5076141357422, + 130.17654418945312, + 129.22286987304688, + 128.9844512939453, + 129.93812561035156, + 128.03077697753906, + 129.46128845214844, + 128.9844512939453 + ] + ], + "median_us": 130.1288604736328, + "tflops": 330.0549378798458 + }, + "staged": { + "round_mean_us": [ + 127.1963119506836, + 126.55258178710936, + 127.60162353515626 + ], + "launch_us": [ + [ + 125.40817260742188, + 128.03077697753906, + 126.12342834472656, + 127.55393981933594, + 127.55393981933594, + 127.7923583984375, + 129.22286987304688, + 126.83868408203125, + 127.07710266113281, + 126.36184692382812 + ], + [ + 126.83868408203125, + 127.55393981933594, + 126.83868408203125, + 124.21607971191406, + 126.36184692382812, + 128.03077697753906, + 126.12342834472656, + 125.40817260742188, + 127.31552124023438, + 126.83868408203125 + ], + [ + 126.60026550292969, + 128.5076141357422, + 127.7923583984375, + 127.55393981933594, + 127.7923583984375, + 127.31552124023438, + 128.03077697753906, + 126.83868408203125, + 128.03077697753906, + 127.55393981933594 + ] + ], + "median_us": 127.1963119506836, + "tflops": 337.6644519115648 + } + } + }, + { + "id": 6, + "b": 256, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 248.0030059814453, + 247.12085723876953, + 247.1923828125 + ], + "launch_us": [ + [ + 246.7632293701172, + 247.00164794921875, + 245.80955505371094, + 252.00843811035156, + 247.2400665283203, + 249.1474151611328, + 251.29318237304688, + 246.52481079101562, + 247.2400665283203, + 247.00164794921875 + ], + [ + 246.28639221191406, + 248.43215942382812, + 247.00164794921875, + 248.19374084472656, + 246.52481079101562, + 248.19374084472656, + 247.00164794921875, + 247.00164794921875, + 245.57113647460938, + 247.00164794921875 + ], + [ + 246.7632293701172, + 246.28639221191406, + 247.2400665283203, + 247.00164794921875, + 247.955322265625, + 246.52481079101562, + 246.28639221191406, + 248.6705780029297, + 248.19374084472656, + 247.00164794921875 + ] + ], + "median_us": 247.1923828125, + "tflops": 347.49997124772347 + }, + "manual": { + "round_mean_us": [ + 248.93283843994138, + 249.45735931396484, + 249.05204772949216 + ], + "launch_us": [ + [ + 248.6705780029297, + 251.0547637939453, + 248.19374084472656, + 249.62425231933594, + 247.71690368652344, + 248.43215942382812, + 250.33950805664062, + 247.71690368652344, + 248.90899658203125, + 248.6705780029297 + ], + [ + 248.43215942382812, + 249.8626708984375, + 250.10108947753906, + 250.10108947753906, + 247.71690368652344, + 247.71690368652344, + 248.90899658203125, + 251.29318237304688, + 250.33950805664062, + 250.10108947753906 + ], + [ + 252.00843811035156, + 247.47848510742188, + 249.38583374023438, + 249.38583374023438, + 248.90899658203125, + 247.71690368652344, + 250.81634521484375, + 249.38583374023438, + 247.955322265625, + 247.47848510742188 + ] + ], + "median_us": 249.05204772949216, + "tflops": 344.9051983435188 + }, + "staged": { + "round_mean_us": [ + 245.07045745849612, + 245.40424346923828, + 245.47576904296872 + ], + "launch_us": [ + [ + 244.140625, + 244.61746215820312, + 245.09429931640625, + 245.57113647460938, + 243.90220642089844, + 245.09429931640625, + 246.0479736328125, + 245.3327178955078, + 246.0479736328125, + 244.8558807373047 + ], + [ + 249.38583374023438, + 246.0479736328125, + 243.18695068359375, + 244.140625, + 245.09429931640625, + 244.8558807373047, + 244.61746215820312, + 247.00164794921875, + 245.3327178955078, + 244.37904357910156 + ], + [ + 245.3327178955078, + 244.140625, + 245.80955505371094, + 247.00164794921875, + 245.3327178955078, + 247.00164794921875, + 246.0479736328125, + 244.37904357910156, + 244.140625, + 245.57113647460938 + ] + ], + "median_us": 245.40424346923828, + "tflops": 350.0320316619447 + } + } + }, + { + "id": 7, + "b": 512, + "sq": 4, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 490.570068359375, + 489.7832870483399, + 490.117073059082 + ], + "launch_us": [ + [ + 491.85752868652344, + 490.42701721191406, + 490.9038543701172, + 490.6654357910156, + 490.6654357910156, + 490.42701721191406, + 488.51966857910156, + 491.3806915283203, + 491.3806915283203, + 489.4733428955078 + ], + [ + 491.6191101074219, + 495.1953887939453, + 489.95018005371094, + 488.7580871582031, + 488.28125, + 489.4733428955078, + 491.85752868652344, + 486.8507385253906, + 487.8044128417969, + 488.04283142089844 + ], + [ + 489.4733428955078, + 489.23492431640625, + 491.14227294921875, + 490.9038543701172, + 490.1885986328125, + 491.3806915283203, + 492.095947265625, + 488.51966857910156, + 490.42701721191406, + 487.8044128417969 + ] + ], + "median_us": 490.117073059082, + "tflops": 350.5258259372863 + }, + "manual": { + "round_mean_us": [ + 495.86296081542974, + 495.55301666259766, + 495.88680267333984 + ], + "launch_us": [ + [ + 496.14906311035156, + 494.24171447753906, + 494.0032958984375, + 494.7185516357422, + 494.7185516357422, + 499.48692321777344, + 495.4338073730469, + 499.725341796875, + 493.04962158203125, + 497.1027374267578 + ], + [ + 496.86431884765625, + 495.67222595214844, + 494.4801330566406, + 495.91064453125, + 497.3411560058594, + 494.24171447753906, + 497.1027374267578, + 495.67222595214844, + 493.2880401611328, + 494.95697021484375 + ], + [ + 495.67222595214844, + 495.91064453125, + 497.8179931640625, + 493.76487731933594, + 496.14906311035156, + 496.3874816894531, + 494.95697021484375, + 496.14906311035156, + 495.1953887939453, + 496.86431884765625 + ] + ], + "median_us": 495.86296081542974, + "tflops": 346.4640544183476 + }, + "staged": { + "round_mean_us": [ + 479.4597625732422, + 479.0782928466797, + 479.4120788574219 + ], + "launch_us": [ + [ + 480.65185546875, + 481.3671112060547, + 478.0292510986328, + 477.79083251953125, + 480.41343688964844, + 479.2213439941406, + 480.65185546875, + 478.2676696777344, + 479.69818115234375, + 478.50608825683594 + ], + [ + 477.3139953613281, + 478.50608825683594, + 479.2213439941406, + 478.50608825683594, + 480.65185546875, + 479.69818115234375, + 478.0292510986328, + 478.2676696777344, + 477.5524139404297, + 483.0360412597656 + ], + [ + 477.79083251953125, + 477.79083251953125, + 481.3671112060547, + 478.0292510986328, + 479.69818115234375, + 479.2213439941406, + 481.3671112060547, + 481.1286926269531, + 479.69818115234375, + 478.0292510986328 + ] + ], + "median_us": 479.4120788574219, + "tflops": 358.3528647201509 + } + } + }, + { + "id": 8, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp8", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 1045.9661483764648, + 1046.6814041137695, + 1046.156883239746 + ], + "launch_us": [ + [ + 1046.1807250976562, + 1047.372817993164, + 1044.9886322021484, + 1047.372817993164, + 1045.4654693603516, + 1046.6575622558594, + 1044.9886322021484, + 1047.1343994140625, + 1045.4654693603516, + 1044.0349578857422 + ], + [ + 1044.2733764648438, + 1045.7038879394531, + 1045.7038879394531, + 1047.6112365722656, + 1051.9027709960938, + 1046.1807250976562, + 1045.4654693603516, + 1048.3264923095703, + 1045.4654693603516, + 1046.1807250976562 + ], + [ + 1047.6112365722656, + 1046.1807250976562, + 1046.6575622558594, + 1045.22705078125, + 1045.22705078125, + 1046.1807250976562, + 1043.7965393066406, + 1046.4191436767578, + 1046.895980834961, + 1047.372817993164 + ] + ], + "median_us": 1046.156883239746, + "tflops": 328.4377220899653 + }, + "manual": { + "round_mean_us": [ + 1034.0452194213867, + 1034.5458984375, + 1033.8544845581055 + ], + "launch_us": [ + [ + 1032.5908660888672, + 1033.5445404052734, + 1034.9750518798828, + 1035.6903076171875, + 1035.451889038086, + 1030.6835174560547, + 1034.0213775634766, + 1034.2597961425781, + 1035.451889038086, + 1033.782958984375 + ], + [ + 1033.5445404052734, + 1034.7366333007812, + 1034.7366333007812, + 1035.451889038086, + 1036.1671447753906, + 1034.0213775634766, + 1031.8756103515625, + 1034.4982147216797, + 1035.928726196289, + 1034.4982147216797 + ], + [ + 1031.8756103515625, + 1033.0677032470703, + 1035.2134704589844, + 1034.0213775634766, + 1035.928726196289, + 1031.1603546142578, + 1034.9750518798828, + 1034.9750518798828, + 1036.1671447753906, + 1031.1603546142578 + ] + ], + "median_us": 1034.0452194213867, + "tflops": 332.2846788772587 + }, + "staged": { + "round_mean_us": [ + 998.2824325561525, + 997.9009628295898, + 998.0916976928711 + ], + "launch_us": [ + [ + 997.5433349609375, + 998.0201721191406, + 997.0664978027344, + 998.7354278564453, + 999.2122650146484, + 998.0201721191406, + 998.0201721191406, + 998.4970092773438, + 998.7354278564453, + 998.9738464355469 + ], + [ + 998.0201721191406, + 997.0664978027344, + 997.5433349609375, + 997.7817535400391, + 998.0201721191406, + 998.2585906982422, + 997.5433349609375, + 997.3049163818359, + 998.4970092773438, + 998.9738464355469 + ], + [ + 997.5433349609375, + 998.2585906982422, + 998.4970092773438, + 999.2122650146484, + 997.7817535400391, + 997.5433349609375, + 998.4970092773438, + 998.0201721191406, + 998.7354278564453, + 996.8280792236328 + ] + ], + "median_us": 998.0916976928711, + "tflops": 344.2543250026416 + } + } + }, + { + "id": 9, + "b": 1, + "sq": 4096, + "topk": 128, + "extra_topk": 512, + "extra_format": "fp4", + "correctness": { + "reference": true, + "manual": true, + "staged": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 939.7268295288086, + 934.1716766357422, + 936.1028671264648 + ], + "launch_us": [ + [ + 938.8923645019531, + 950.0980377197266, + 935.7929229736328, + 946.9985961914062, + 934.3624114990234, + 938.8923645019531, + 936.0313415527344, + 947.2370147705078, + 935.5545043945312, + 933.4087371826172 + ], + [ + 935.3160858154297, + 934.3624114990234, + 935.5545043945312, + 933.1703186035156, + 932.6934814453125, + 934.3624114990234, + 934.1239929199219, + 934.1239929199219, + 934.1239929199219, + 933.8855743408203 + ], + [ + 932.9319000244141, + 949.3827819824219, + 943.1838989257812, + 932.9319000244141, + 934.600830078125, + 932.2166442871094, + 932.2166442871094, + 934.600830078125, + 934.600830078125, + 934.3624114990234 + ] + ], + "median_us": 936.1028671264648, + "tflops": 367.0508827034507 + }, + "manual": { + "round_mean_us": [ + 953.1497955322266, + 953.4597396850586, + 953.4597396850586 + ], + "launch_us": [ + [ + 954.1511535644531, + 952.0053863525391, + 952.4822235107422, + 953.67431640625, + 953.4358978271484, + 954.8664093017578, + 953.1974792480469, + 951.7669677734375, + 952.9590606689453, + 952.9590606689453 + ], + [ + 953.4358978271484, + 952.9590606689453, + 952.7206420898438, + 953.9127349853516, + 953.67431640625, + 953.1974792480469, + 952.9590606689453, + 954.6279907226562, + 953.4358978271484, + 953.67431640625 + ], + [ + 953.67431640625, + 952.9590606689453, + 952.4822235107422, + 954.1511535644531, + 952.7206420898438, + 955.1048278808594, + 952.2438049316406, + 952.7206420898438, + 954.3895721435547, + 954.1511535644531 + ] + ], + "median_us": 953.4597396850586, + "tflops": 360.3690532266157 + }, + "staged": { + "round_mean_us": [ + 957.6082229614258, + 957.2744369506836, + 957.8943252563477 + ], + "launch_us": [ + [ + 957.9658508300781, + 957.2505950927734, + 957.489013671875, + 957.7274322509766, + 957.2505950927734, + 957.2505950927734, + 956.7737579345703, + 960.5884552001953, + 956.7737579345703, + 957.0121765136719 + ], + [ + 957.9658508300781, + 957.489013671875, + 957.2505950927734, + 956.7737579345703, + 957.0121765136719, + 957.489013671875, + 957.489013671875, + 957.9658508300781, + 956.5353393554688, + 956.7737579345703 + ], + [ + 957.0121765136719, + 957.489013671875, + 961.7805480957031, + 957.9658508300781, + 956.2969207763672, + 957.9658508300781, + 958.2042694091797, + 957.0121765136719, + 957.2505950927734, + 957.9658508300781 + ] + ], + "median_us": 957.6082229614258, + "tflops": 358.80788765307074 + } + } + } + ], + "provenance": { + "head": "00e7b9bb64c892c5323f42bede764e8940ceb4cd", + "upstream_main": "876a59e6eb8091ea608c14a84d5300190c3735ce", + "reference_git_head": "a123d0b0191e0da7aa1e044f8644e0989cc63220", + "source_files": { + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/manual.py": "0de7427f4514b84a58c4cf6f789dc0fba3f4f1617ebe05cebc9a08675c43abbb", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/decode.py": "394e9d46a9478723eb512eaa66be04d2d719f0fb0df9826be98537da3d904b3a", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/run_decode.py": "80d4396d3422c618b92475d1c2be3418226372a0a6b9044757842970065296db", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/core.py": "1cd425ce94f37dee8d89e4b25b7c882da0c3b899df2f098db3276eac6daf4813", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/staged.py": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/quantization.py": "4ccd402d64f8e45ae362ef36868964baaae3a187769659e913912c7b701ae064", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/test_staged.py": "a5428baae9dfbd14de0f63664652c1ab5bef647474523c16746a76e0d9968c4e", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/run.py": "df9df8607181c62760c8cee0474d4c015e1d892bee0623543a663e12a18b7e6b", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/manual_intrinsics.h": "4f20b6d3e05ceeeaee3f35e739c7bcfcb6e3ef5bbdbba2f26f33d6f236288049", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/.tlcli/flashmla-phase-perf/decode-boundary-check/run_decode_boundaries.py": "8d4580b4d1bfed6c76a42d5e319b57ecb38847743f33f7e893a965200f713bd4", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/lib.py": "985306103578f09b7561dfa89acd8035c276835396af7bfe0e926a465520c2bc", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/generate.py": "691828edb97cfa05aabe7e3c1091e4d23dd9435e8816568b348f50379b268d08", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/utils.py": "ca17b58340f85d3e64c53b21cbf241e7fe61c62ec6359bcc16316fc05a26afcf", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/quant.py": "2bc2bbe8a923e2dfd2a52a3ba8c25370e13cd1361c31a0ed2a8b75f769f7b17e", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/ref.py": "696b87713465a958ec9540903e8a4876106f28638467468fa0070c96182de4f6", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/bench.py": "2a707bc8a647e19593d6691b7d62446eea605bc467f8323db2ca2e8fdbcbebca", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/flash_mla/npu.cpython-312-aarch64-linux-gnu.so": "70d7060d8e59ccb56e402a20e63cc0c9d8b91ffe6b32884afd26365f2c391fda" + }, + "native_libraries": { + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtilelang.so": "489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtvm_compiler.so": "572a2285726e0d0565f47d4e56d757d31c956959605ff126d4725f8a510004a4", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtvm_runtime.so": "2b6cd5655295f8cd5915abcb39929aa4f0a2522579d2a6440c5503820e958b00" + }, + "method": { + "rounds": 3, + "repeat": 10, + "warmup_per_round": 10, + "flush_l2_bytes": 8000000000, + "order": "reference/manual/staged, reversed on odd rounds", + "summary": "median of 3 round means", + "profiler": "kernelkit / torch_npu msprof" + }, + "prepared_kernels": [ + { + "implementation": "manual", + "case": "perf-decode-0", + "arguments": { + "nq": 256, + "nk": 327680, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 34603008, + "extra_storage_bytes": 138412032, + "cache_hint": 0 + }, + "asc_sha256": "e186b8d2f7e071f2c5fdf9b723e4f9f0c8e4fc526e3024e96cbeabf24fdf033a", + "binary_sha256": "13309e59cd9ecc4aacb3957873774dc9f79bfc46f909a18ee6936b1c09064a8c", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.49945280099928 + }, + { + "implementation": "staged", + "case": "perf-decode-0", + "arguments": { + "nq": 256, + "nk": 327680, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 34603008, + "extra_storage_bytes": 138412032, + "cache_hint": 0 + }, + "asc_sha256": "18529e510169b4530bedd5507ba20f4bd763d866ec8389596f6883fc4cc86170", + "binary_sha256": "4b41d1cee8679af64d1ce4079716e9c42a9aad6f49d07e078eb4003a05a0be63", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 206848, + "total_bytes": 206848, + "limit_bytes": 253952, + "remaining_bytes": 47104 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 343.57139671999903 + }, + { + "implementation": "manual", + "case": "perf-decode-1", + "arguments": { + "nq": 512, + "nk": 851968, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 69206016, + "extra_storage_bytes": 380633088, + "cache_hint": 0 + }, + "asc_sha256": "f8bf4d87875c4a908bc4c8f66b9d32aa54f31a08e83ec16068c385765b880825", + "binary_sha256": "58fe53c9c31517e638c63b6f7803fb60ff16621162b8993fdcb124affd6588cd", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.580423954001162 + }, + { + "implementation": "staged", + "case": "perf-decode-1", + "arguments": { + "nq": 512, + "nk": 851968, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 69206016, + "extra_storage_bytes": 380633088, + "cache_hint": 0 + }, + "asc_sha256": "3ad1d1071aa774f00824e62f8d04badca2bb3744e15d625d7f761b1840e5f1d0", + "binary_sha256": "95bd678b0c9688d21e4ccd12caf4d5db88cd0477a51126e35fe6d05f6f87f170", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 206848, + "total_bytes": 206848, + "limit_bytes": 253952, + "remaining_bytes": 47104 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 342.85152793499583 + }, + { + "implementation": "manual", + "case": "perf-decode-2", + "arguments": { + "nq": 1024, + "nk": 1703936, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 138412032, + "extra_storage_bytes": 761266176, + "cache_hint": 0 + }, + "asc_sha256": "5296ed8009f14208ea912fbeefcd56d570e829c9ea9c635701f68bf99adcea54", + "binary_sha256": "82fc3e72af9d40ca3717462e744b73462eb17ac0d15da2bf7d1eda9dadddc278", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.632557354998426 + }, + { + "implementation": "staged", + "case": "perf-decode-2", + "arguments": { + "nq": 1024, + "nk": 1703936, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 138412032, + "extra_storage_bytes": 761266176, + "cache_hint": 0 + }, + "asc_sha256": "73b717ff6c9806c66b744c9814f88a93b1fd89dd5eeeaecc34be33aca9f19137", + "binary_sha256": "9932a06d37bbe16eba50c6c81048771ab07c62e33dbb489c2eba81103195772b", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 206848, + "total_bytes": 206848, + "limit_bytes": 253952, + "remaining_bytes": 47104 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 345.9321977310028 + }, + { + "implementation": "manual", + "case": "perf-decode-3", + "arguments": { + "nq": 2048, + "nk": 3407872, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 276824064, + "extra_storage_bytes": 1522532352, + "cache_hint": 0 + }, + "asc_sha256": "910fa21189fc861adf95490354fcf6936d5ffeb6ee131c42e75d915ca6926996", + "binary_sha256": "b8dc296ab45b9f694262c822840b659a95c088cfba0bc69f16dc60b86746b4d6", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.398867405005149 + }, + { + "implementation": "staged", + "case": "perf-decode-3", + "arguments": { + "nq": 2048, + "nk": 3407872, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 276824064, + "extra_storage_bytes": 1522532352, + "cache_hint": 0 + }, + "asc_sha256": "4d0eb115fecfc018daeda0ad5c600e01f545c0032f9012825b5c71b1e944981c", + "binary_sha256": "9f5b7fb70c7aaadfd83cdc068a0fb7069497e489c134dfe3a4b34f021c31759b", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 206848, + "total_bytes": 206848, + "limit_bytes": 253952, + "remaining_bytes": 47104 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 344.6234449600015 + }, + { + "implementation": "manual", + "case": "perf-decode-4", + "arguments": { + "nq": 256, + "nk": 327680, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 34603008, + "extra_storage_bytes": 75497472, + "cache_hint": 0 + }, + "asc_sha256": "040c98bc726efb5d9ddae2dba8907a224f80630406c8fe608aedc46ca4f5476d", + "binary_sha256": "bb040c3ec96af0427f93c8573584fa80e1f7974411581fcf164dbea1d6749aba", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 11.550360281005851 + }, + { + "implementation": "staged", + "case": "perf-decode-4", + "arguments": { + "nq": 256, + "nk": 327680, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 34603008, + "extra_storage_bytes": 75497472, + "cache_hint": 0 + }, + "asc_sha256": "1be5a9d54cbdd1bacf994691572c48feb0b004e64bfbeca4732ab5f2153983ba", + "binary_sha256": "18927d6db596085aa2862d5707ab92835b193c25736020b47be1477be9a18b3b", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 234496, + "total_bytes": 234496, + "limit_bytes": 253952, + "remaining_bytes": 19456 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 317.788200918003 + }, + { + "implementation": "manual", + "case": "perf-decode-5", + "arguments": { + "nq": 512, + "nk": 851968, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 69206016, + "extra_storage_bytes": 207618048, + "cache_hint": 0 + }, + "asc_sha256": "9233cfacd2c475ec02e114f516a5ca0b9a33d31c578f2831c28c3cc05071e9ba", + "binary_sha256": "7ee58ce76ff9a4415565fdf0ed370d95e180528f19325e878ba55935b10be6e1", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 11.080378612998175 + }, + { + "implementation": "staged", + "case": "perf-decode-5", + "arguments": { + "nq": 512, + "nk": 851968, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 69206016, + "extra_storage_bytes": 207618048, + "cache_hint": 0 + }, + "asc_sha256": "fe1d147342751782e431be9433c1b419b0cb1b6ca99128e6df07adbbd56d7d04", + "binary_sha256": "557f52ce86f7e6011f475a740665d230d7b625b9568b0e3721e26c831481fcc7", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 234496, + "total_bytes": 234496, + "limit_bytes": 253952, + "remaining_bytes": 19456 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 324.8401765979943 + }, + { + "implementation": "manual", + "case": "perf-decode-6", + "arguments": { + "nq": 1024, + "nk": 1703936, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 138412032, + "extra_storage_bytes": 415236096, + "cache_hint": 0 + }, + "asc_sha256": "3aeaa1805201358cfb3a132f071dd7df7493a59ef333518e321b5a74c655de03", + "binary_sha256": "402d920fa42d979691a1787bd546b2e2e2b6bf56c0e30184b884607b8a06ad1d", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.571845121012302 + }, + { + "implementation": "staged", + "case": "perf-decode-6", + "arguments": { + "nq": 1024, + "nk": 1703936, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 138412032, + "extra_storage_bytes": 415236096, + "cache_hint": 0 + }, + "asc_sha256": "e653bb3f922a62847a33d0714d19cef7e09ccc1c7511942f2bc79edb7f11889f", + "binary_sha256": "f9a1e5cf40401afb20bfc4a0b7baf1f819b4b7e38eed6a4d0958c5b807df7b3d", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 234496, + "total_bytes": 234496, + "limit_bytes": 253952, + "remaining_bytes": 19456 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 321.8030986609956 + }, + { + "implementation": "manual", + "case": "perf-decode-7", + "arguments": { + "nq": 2048, + "nk": 3407872, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 276824064, + "extra_storage_bytes": 830472192, + "cache_hint": 0 + }, + "asc_sha256": "b0ea1e00e3ed118c8e4bae4a9b950b7e445b109d1963844727fe76f4c8ef713a", + "binary_sha256": "e2c5f9dde473247a1e4f3a3e4719145b5542118da79789e23920e49a20a26694", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.498413045992493 + }, + { + "implementation": "staged", + "case": "perf-decode-7", + "arguments": { + "nq": 2048, + "nk": 3407872, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 32768, + 512, + 1 + ], + "index_stride": 128, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 512, + "variable_extra_lengths": false, + "kv_storage_bytes": 276824064, + "extra_storage_bytes": 830472192, + "cache_hint": 0 + }, + "asc_sha256": "602573d9b004dc767456880dde4e687b71a90270084567f77914c153679fe567", + "binary_sha256": "b7db84c44ef6fc9541e428b37f2d2f0a03249d18c7f822ad8b08e49090424687", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 234496, + "total_bytes": 234496, + "limit_bytes": 253952, + "remaining_bytes": 19456 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 321.8236621410033 + }, + { + "implementation": "manual", + "case": "perf-decode-8", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 256, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 640, + "variable_extra_lengths": false, + "kv_storage_bytes": 2162688, + "extra_storage_bytes": 2162688, + "cache_hint": 1 + }, + "asc_sha256": "f781a501006693cceed7f43826b43f2155c6ad42302add410a56f504473fe78f", + "binary_sha256": "497080f417907009a334a9ee67d37b3e479af0ba8d1745d1ba12ac1e5b5dc972", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.710857295998721 + }, + { + "implementation": "staged", + "case": "perf-decode-8", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 256, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp8", + "extra_index_stride": 640, + "variable_extra_lengths": false, + "kv_storage_bytes": 2162688, + "extra_storage_bytes": 2162688, + "cache_hint": 1 + }, + "asc_sha256": "2d8aaf6adf2d1f5fab53344ef3aeef35bc23c9b17366f5af70ddef5d7d26b1d7", + "binary_sha256": "80a65de42232cb6bc3ba27f8544470f8c39708268c5a1052b0e5ca301c3d0fd0", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 206848, + "total_bytes": 206848, + "limit_bytes": 253952, + "remaining_bytes": 47104 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 345.6519301980006 + }, + { + "implementation": "manual", + "case": "perf-decode-9", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 256, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 640, + "variable_extra_lengths": false, + "kv_storage_bytes": 2162688, + "extra_storage_bytes": 1179648, + "cache_hint": 1 + }, + "asc_sha256": "c928c849d87041e05f3b9dc0dc5d71c80405f4daae7de161d0cef3de704578c1", + "binary_sha256": "91e61f22fe8a85d625c204ca6ea5c3d6f9cd70c93ba148d387aa0474f552a116", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 241920, + "total_bytes": 241920, + "limit_bytes": 253952, + "remaining_bytes": 12032 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.515198100009002 + }, + { + "implementation": "staged", + "case": "perf-decode-9", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 128, + "sink": true, + "variable_lengths": false, + "scale": 0.032352028870043036, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 256, + "kv_format": "fp8", + "extra_topk": 512, + "extra_format": "fp4", + "extra_index_stride": 640, + "variable_extra_lengths": false, + "kv_storage_bytes": 2162688, + "extra_storage_bytes": 1179648, + "cache_hint": 1 + }, + "asc_sha256": "99996bc1cf7b6c364efa78fb0d7c2205b457526536424d5e96bbb7cb4d8cf6a7", + "binary_sha256": "ee7a02935a2f0fd218da79024c7fa7bff12a1aadc8d3e7c0af2b9b6e7b6ab1f7", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 234496, + "total_bytes": 234496, + "limit_bytes": 253952, + "remaining_bytes": 19456 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 321.8354039859987 + } + ], + "compiler": "CANN 9.2.0 / Bisheng 15.0.5, dav-3510", + "compile_flags": [ + "-Ofast", + "-mllvm", + "-enable-hiipu-vf-loop-unroll" + ], + "cache_policy": "Fresh isolated cache populated by this compiler before timing; source/library/file hashes verified and new compilation forbidden during measurement.", + "validation_scope": "Output comparisons at all benchmark shapes against the pinned FlashMLA implementation; not a rerun of the separate correctness matrix.", + "artifacts": "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/.tlcli/flashmla-phase-perf/main-rebase-perf-20261020", + "raw_result_sha256": "fbafd4b8f76c7c6a78cc7d040deb19a99be01bfe997e0ef6000aeb87cdecb061", + "audit_sha256": "c25d097d6acc16935e5080f4c87c7307a88f18c8995fc3252487b8e30b8bf676" + } +} diff --git a/examples/ascend/flashmla/results/staged-main-rebase-prefill-bench-2026-10-20.json b/examples/ascend/flashmla/results/staged-main-rebase-prefill-bench-2026-10-20.json new file mode 100644 index 0000000000..a9719095b0 --- /dev/null +++ b/examples/ascend/flashmla/results/staged-main-rebase-prefill-bench-2026-10-20.json @@ -0,0 +1,947 @@ +{ + "torch": "2.13.0+cpu", + "torch_npu": "2.13.0rc1", + "device": "Ascend950DT_9581", + "visible_devices": "6", + "cases": [ + { + "nq": 4096, + "nk": 4096, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true, + "staged_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 853.3239364624023, + 854.5160293579102, + 853.1570434570312 + ], + "launch_us": [ + [ + 846.6243743896484, + 857.5916290283203, + 853.3000946044922, + 850.2006530761719, + 864.2673492431641, + 851.1543273925781, + 852.3464202880859, + 851.1543273925781, + 855.9226989746094, + 850.677490234375 + ], + [ + 866.6515350341797, + 860.4526519775391, + 854.7306060791016, + 849.2469787597656, + 845.9091186523438, + 855.2074432373047, + 849.0085601806641, + 849.0085601806641, + 866.8899536132812, + 848.0548858642578 + ], + [ + 847.3396301269531, + 863.7905120849609, + 855.9226989746094, + 861.1679077148438, + 849.4853973388672, + 852.5848388671875, + 852.1080017089844, + 849.7238159179688, + 849.9622344970703, + 849.4853973388672 + ] + ], + "median_us": 853.3239364624023, + "tflops": 402.6576180488276 + }, + "manual": { + "round_mean_us": [ + 848.5317230224609, + 848.6032485961914, + 849.0800857543945 + ], + "launch_us": [ + [ + 847.8164672851562, + 848.5317230224609, + 847.8164672851562, + 847.8164672851562, + 849.0085601806641, + 847.8164672851562, + 849.0085601806641, + 849.4853973388672, + 849.2469787597656, + 848.7701416015625 + ], + [ + 849.2469787597656, + 848.2933044433594, + 849.0085601806641, + 847.3396301269531, + 848.5317230224609, + 849.9622344970703, + 848.2933044433594, + 848.0548858642578, + 848.0548858642578, + 849.2469787597656 + ], + [ + 848.2933044433594, + 848.5317230224609, + 848.7701416015625, + 848.7701416015625, + 848.2933044433594, + 848.0548858642578, + 849.2469787597656, + 853.5385131835938, + 848.2933044433594, + 849.0085601806641 + ] + ], + "median_us": 848.6032485961914, + "tflops": 404.8975587218157 + }, + "staged": { + "round_mean_us": [ + 864.2435073852539, + 864.863395690918, + 864.1719818115234 + ], + "launch_us": [ + [ + 864.0289306640625, + 864.2673492431641, + 864.0289306640625, + 865.4594421386719, + 864.0289306640625, + 864.5057678222656, + 864.5057678222656, + 865.4594421386719, + 863.0752563476562, + 863.0752563476562 + ], + [ + 864.5057678222656, + 864.2673492431641, + 863.7905120849609, + 866.8899536132812, + 866.1746978759766, + 865.4594421386719, + 864.5057678222656, + 864.5057678222656, + 863.3136749267578, + 865.2210235595703 + ], + [ + 863.7905120849609, + 863.0752563476562, + 864.7441864013672, + 864.0289306640625, + 863.5520935058594, + 865.4594421386719, + 863.0752563476562, + 863.7905120849609, + 865.4594421386719, + 864.7441864013672 + ] + ], + "median_us": 864.2435073852539, + "tflops": 397.5701069708292 + } + } + }, + { + "nq": 4096, + "nk": 8192, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true, + "staged_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 856.3041687011719, + 854.8736572265625, + 852.7278900146484 + ], + "launch_us": [ + [ + 849.4853973388672, + 861.6447448730469, + 852.5848388671875, + 865.936279296875, + 851.8695831298828, + 849.9622344970703, + 849.9622344970703, + 853.0616760253906, + 866.6515350341797, + 861.8831634521484 + ], + [ + 856.3995361328125, + 850.2006530761719, + 861.8831634521484, + 850.9159088134766, + 853.7769317626953, + 867.1283721923828, + 848.7701416015625, + 853.7769317626953, + 855.2074432373047, + 850.677490234375 + ], + [ + 852.1080017089844, + 851.8695831298828, + 850.4390716552734, + 849.9622344970703, + 866.4131164550781, + 853.3000946044922, + 849.9622344970703, + 853.0616760253906, + 851.1543273925781, + 849.0085601806641 + ] + ], + "median_us": 854.8736572265625, + "tflops": 401.92767758772834 + }, + "manual": { + "round_mean_us": [ + 850.4390716552734, + 850.0099182128906, + 850.4629135131836 + ], + "launch_us": [ + [ + 850.4390716552734, + 850.4390716552734, + 850.4390716552734, + 848.7701416015625, + 850.677490234375, + 853.5385131835938, + 850.9159088134766, + 851.1543273925781, + 849.4853973388672, + 848.5317230224609 + ], + [ + 849.7238159179688, + 849.2469787597656, + 850.4390716552734, + 849.7238159179688, + 849.7238159179688, + 851.1543273925781, + 849.4853973388672, + 849.4853973388672, + 850.4390716552734, + 850.677490234375 + ], + [ + 849.7238159179688, + 849.9622344970703, + 850.9159088134766, + 850.677490234375, + 850.9159088134766, + 850.2006530761719, + 848.7701416015625, + 849.4853973388672, + 851.3927459716797, + 852.5848388671875 + ] + ], + "median_us": 850.4390716552734, + "tflops": 404.02351577195367 + }, + "staged": { + "round_mean_us": [ + 866.6753768920898, + 866.4369583129883, + 866.8899536132812 + ], + "launch_us": [ + [ + 866.1746978759766, + 866.6515350341797, + 865.6978607177734, + 866.4131164550781, + 865.4594421386719, + 868.5588836669922, + 868.0820465087891, + 867.3667907714844, + 866.4131164550781, + 865.936279296875 + ], + [ + 866.1746978759766, + 866.6515350341797, + 866.1746978759766, + 867.1283721923828, + 867.3667907714844, + 866.6515350341797, + 866.4131164550781, + 865.4594421386719, + 865.936279296875, + 866.4131164550781 + ], + [ + 866.8899536132812, + 866.8899536132812, + 866.1746978759766, + 867.8436279296875, + 866.1746978759766, + 868.0820465087891, + 865.936279296875, + 867.6052093505859, + 865.2210235595703, + 868.0820465087891 + ] + ], + "median_us": 866.6753768920898, + "tflops": 396.4545351595716 + } + } + }, + { + "nq": 4096, + "nk": 32768, + "topk": 640, + "sink": true, + "lengths": false, + "amplified": false, + "all_invalid_rows": false, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_strides": [ + 1536, + 768, + 1 + ], + "correctness": { + "manual_vs_reference": true, + "staged_vs_reference": true + }, + "performance": { + "reference": { + "round_mean_us": [ + 861.6447448730469, + 862.0500564575195, + 861.9546890258789 + ], + "launch_us": [ + [ + 861.6447448730469, + 861.1679077148438, + 861.6447448730469, + 861.1679077148438, + 862.8368377685547, + 861.6447448730469, + 861.8831634521484, + 861.1679077148438, + 861.8831634521484, + 861.4063262939453 + ], + [ + 862.12158203125, + 861.6447448730469, + 861.8831634521484, + 861.4063262939453, + 862.3600006103516, + 861.4063262939453, + 862.12158203125, + 862.5984191894531, + 862.5984191894531, + 862.3600006103516 + ], + [ + 862.12158203125, + 861.4063262939453, + 862.3600006103516, + 860.9294891357422, + 862.5984191894531, + 862.3600006103516, + 861.8831634521484, + 860.9294891357422, + 862.3600006103516, + 862.5984191894531 + ] + ], + "median_us": 861.9546890258789, + "tflops": 398.6258071967911 + }, + "manual": { + "round_mean_us": [ + 861.7162704467773, + 861.811637878418, + 861.3348007202148 + ], + "launch_us": [ + [ + 862.8368377685547, + 859.7373962402344, + 861.8831634521484, + 861.6447448730469, + 860.4526519775391, + 861.4063262939453, + 862.5984191894531, + 862.12158203125, + 862.3600006103516, + 862.12158203125 + ], + [ + 860.4526519775391, + 861.4063262939453, + 859.9758148193359, + 862.3600006103516, + 862.12158203125, + 862.3600006103516, + 861.4063262939453, + 863.5520935058594, + 862.3600006103516, + 862.12158203125 + ], + [ + 859.4989776611328, + 859.2605590820312, + 862.8368377685547, + 862.5984191894531, + 862.3600006103516, + 863.0752563476562, + 862.5984191894531, + 859.7373962402344, + 859.4989776611328, + 861.8831634521484 + ] + ], + "median_us": 861.7162704467773, + "tflops": 398.7360984861685 + }, + "staged": { + "round_mean_us": [ + 873.9709854125977, + 874.0425109863281, + 873.7802505493164 + ], + "launch_us": [ + [ + 874.2809295654297, + 872.8504180908203, + 874.7577667236328, + 873.0888366699219, + 874.0425109863281, + 873.8040924072266, + 873.3272552490234, + 874.5193481445312, + 874.7577667236328, + 874.2809295654297 + ], + [ + 873.565673828125, + 874.2809295654297, + 875.7114410400391, + 873.3272552490234, + 872.6119995117188, + 873.8040924072266, + 874.0425109863281, + 876.6651153564453, + 873.565673828125, + 872.8504180908203 + ], + [ + 873.8040924072266, + 873.565673828125, + 873.3272552490234, + 874.7577667236328, + 874.2809295654297, + 872.1351623535156, + 874.9961853027344, + 873.3272552490234, + 873.3272552490234, + 874.2809295654297 + ] + ], + "median_us": 873.9709854125977, + "tflops": 393.1450693615295 + } + } + } + ], + "provenance": { + "head": "00e7b9bb64c892c5323f42bede764e8940ceb4cd", + "upstream_main": "876a59e6eb8091ea608c14a84d5300190c3735ce", + "reference_git_head": "a123d0b0191e0da7aa1e044f8644e0989cc63220", + "source_files": { + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/manual.py": "0de7427f4514b84a58c4cf6f789dc0fba3f4f1617ebe05cebc9a08675c43abbb", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/decode.py": "394e9d46a9478723eb512eaa66be04d2d719f0fb0df9826be98537da3d904b3a", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/run_decode.py": "80d4396d3422c618b92475d1c2be3418226372a0a6b9044757842970065296db", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/core.py": "1cd425ce94f37dee8d89e4b25b7c882da0c3b899df2f098db3276eac6daf4813", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/staged.py": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/quantization.py": "4ccd402d64f8e45ae362ef36868964baaae3a187769659e913912c7b701ae064", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/test_staged.py": "a5428baae9dfbd14de0f63664652c1ab5bef647474523c16746a76e0d9968c4e", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/run.py": "df9df8607181c62760c8cee0474d4c015e1d892bee0623543a663e12a18b7e6b", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/examples/ascend/flashmla/manual_intrinsics.h": "4f20b6d3e05ceeeaee3f35e739c7bcfcb6e3ef5bbdbba2f26f33d6f236288049", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/.tlcli/flashmla-phase-perf/decode-boundary-check/run_decode_boundaries.py": "8d4580b4d1bfed6c76a42d5e319b57ecb38847743f33f7e893a965200f713bd4", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/lib.py": "985306103578f09b7561dfa89acd8035c276835396af7bfe0e926a465520c2bc", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/generate.py": "691828edb97cfa05aabe7e3c1091e4d23dd9435e8816568b348f50379b268d08", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/utils.py": "ca17b58340f85d3e64c53b21cbf241e7fe61c62ec6359bcc16316fc05a26afcf", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/quant.py": "2bc2bbe8a923e2dfd2a52a3ba8c25370e13cd1361c31a0ed2a8b75f769f7b17e", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/ref.py": "696b87713465a958ec9540903e8a4876106f28638467468fa0070c96182de4f6", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/tests/kernelkit/bench.py": "2a707bc8a647e19593d6691b7d62446eea605bc467f8323db2ca2e8fdbcbebca", + "/weka-hg/prod/deepseek/permanent/jintian/private/tilelang/.tlcli/flashmla-perf/reference/flash_mla/npu.cpython-312-aarch64-linux-gnu.so": "70d7060d8e59ccb56e402a20e63cc0c9d8b91ffe6b32884afd26365f2c391fda" + }, + "native_libraries": { + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtilelang.so": "489c5ace819791fec4a2058c93c0ee9a50fddde505cad2a172f0929619ddda3b", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtvm_compiler.so": "572a2285726e0d0565f47d4e56d757d31c956959605ff126d4725f8a510004a4", + "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/build/lib/libtvm_runtime.so": "2b6cd5655295f8cd5915abcb39929aa4f0a2522579d2a6440c5503820e958b00" + }, + "method": { + "rounds": 3, + "repeat": 10, + "warmup_per_round": 10, + "flush_l2_bytes": 8000000000, + "order": "reference/manual/staged, reversed on odd rounds", + "summary": "median of 3 round means", + "profiler": "kernelkit / torch_npu msprof" + }, + "prepared_kernels": [ + { + "implementation": "manual", + "case": "perf-prefill-4096", + "arguments": { + "nq": 4096, + "nk": 4096, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "cb4459e0391a37bc22ebf85fea8fb4717589b46e12c1296c18bd65658743b21c", + "binary_sha256": "89c0972657ae90b541e800bffa170aeef64cb322dee3054539b863216cdcfdd1", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.099991519004107 + }, + { + "implementation": "staged", + "case": "perf-prefill-4096", + "arguments": { + "nq": 4096, + "nk": 4096, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "19ea8a24defad41cb0a07717308152d7a0651221b330206f9b8da04783745edd", + "binary_sha256": "4e1fa10ed1eb82b2ee48a775f5947e6bdd0f1e5c34d5c5df821bda7ee464ecfd", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 282.1487830469996 + }, + { + "implementation": "manual", + "case": "perf-prefill-8192", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "a58e3fc0ec7eb7854b39c290997fdf66143813e43c1556c220ce5529faa5b09b", + "binary_sha256": "22fdbb2a20485b1f8bc15f9e57e3bd517252a0ba9ccf3aff03bc0457c27d96af", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 10.821810818990343 + }, + { + "implementation": "staged", + "case": "perf-prefill-8192", + "arguments": { + "nq": 4096, + "nk": 8192, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "9683f2afd0718a24ed5665df5d4c227c1cbdbbf29ff7bc82ced19bf22104e12b", + "binary_sha256": "d2a43527343a8e837e142b12ae13630198b8f19852a07a5855bca0e9a18de65f", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 275.82231884299836 + }, + { + "implementation": "manual", + "case": "perf-prefill-32768", + "arguments": { + "nq": 4096, + "nk": 32768, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "e29affb931b627c3848d3a4823f6423d5494ac1150fb0f24269390ad04668633", + "binary_sha256": "0e5dc38f6b724ecec0cf89c74637062c31a1b24c70e55e288f5f49b62dfca4fc", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 524288, + "dynamic_bytes": 0, + "total_bytes": 524288, + "limit_bytes": 524288, + "remaining_bytes": 0 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 147456, + "dynamic_bytes": 0, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 512, + "dynamic_bytes": 0, + "total_bytes": 512, + "limit_bytes": 3072, + "remaining_bytes": 2560 + } + }, + "compile_seconds": 9.834505132006598 + }, + { + "implementation": "staged", + "case": "perf-prefill-32768", + "arguments": { + "nq": 4096, + "nk": 32768, + "topk": 640, + "sink": true, + "variable_lengths": false, + "scale": 0.5, + "q_strides": [ + 41600, + 640, + 1 + ], + "index_stride": 1536 + }, + "asc_sha256": "72045f1e00c9f2124cb5408020c33f889c8bd697d4acb7232d6d4cd1e3ec6733", + "binary_sha256": "6b6e788bc54d74a2155dd348d4383ed2e067962307bafee0c452c9c4a906a4ac", + "resources": { + "UB": { + "static_bytes": 0, + "dynamic_bytes": 253184, + "total_bytes": 253184, + "limit_bytes": 253952, + "remaining_bytes": 768 + }, + "L1": { + "static_bytes": 0, + "dynamic_bytes": 409600, + "total_bytes": 409600, + "limit_bytes": 524288, + "remaining_bytes": 114688 + }, + "L0A": { + "static_bytes": 65536, + "dynamic_bytes": 0, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0B": { + "static_bytes": 0, + "dynamic_bytes": 65536, + "total_bytes": 65536, + "limit_bytes": 65536, + "remaining_bytes": 0 + }, + "L0C": { + "static_bytes": 0, + "dynamic_bytes": 147456, + "total_bytes": 147456, + "limit_bytes": 262144, + "remaining_bytes": 114688 + }, + "SS": { + "static_bytes": 0, + "dynamic_bytes": 64, + "total_bytes": 64, + "limit_bytes": 3072, + "remaining_bytes": 3008 + } + }, + "compile_seconds": 282.12404014199274 + } + ], + "compiler": "CANN 9.2.0 / Bisheng 15.0.5, dav-3510", + "compile_flags": [ + "-Ofast", + "-mllvm", + "-enable-hiipu-vf-loop-unroll" + ], + "cache_policy": "Fresh isolated cache populated by this compiler before timing; source/library/file hashes verified and new compilation forbidden during measurement.", + "validation_scope": "Output comparisons at all benchmark shapes against the pinned FlashMLA implementation; not a rerun of the separate correctness matrix.", + "artifacts": "/weka-hg/prod/deepseek/permanent/jintian/private/tl-cli-asc/workspace/dev-flashmla-manual-example/.tlcli/flashmla-phase-perf/main-rebase-perf-20261020", + "raw_result_sha256": "0c8ff354b9ac8456bf60bbaf01a136df4bd49b7bc0970f02040a35e3b29358d9", + "audit_sha256": "b15059432a53fe6f713937aafe994795430d1325d679d80b93742bc7b6e2f996" + } +} diff --git a/examples/ascend/flashmla/results/staged-manual-order-2026-10-20.json b/examples/ascend/flashmla/results/staged-manual-order-2026-10-20.json new file mode 100644 index 0000000000..f5914e8f03 --- /dev/null +++ b/examples/ascend/flashmla/results/staged-manual-order-2026-10-20.json @@ -0,0 +1,309 @@ +{ + "date": "2026-10-20", + "record_type": "Summary of existing local validation; documentation refresh did not rebuild or rerun tests.", + "implementation_commit": "9c46d1cdab152218a9335fadfcde8816409f402b", + "validation_parent_commit": "e9244819b52faa5775dc777570ff9110d0234ba1", + "validation_revision_note": "The original summary records the parent HEAD before committing the validated changes as implementation_commit. Current committed kernel source and local native-library hashes match that summary.", + "source_file": "examples/ascend/flashmla/staged.py", + "source_sha256": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2", + "native_library": "build/lib/libtilelang.so", + "native_library_sha256": "f63a9bebad45fe496bb7e89c621d29efaf819fca41a61e612648175c65885713", + "kernel_issue_phase_calls": 0, + "protocol": "Independent QK/PV B4 layouts retain the shared ring. Manual per-core/pipe source order gains strict phase separation when ascending stages would reverse it during reconstruction. Gather precedes the delta wait and SS publication.", + "validation": { + "focused_tests": 68, + "autoschedule_tests": 776, + "flashmla_device_tests": 11, + "flashmla_device_result": "11 passed, 106 warnings in 921.15s (0:15:21)", + "device_cases": [ + "Cross-query and repeated-launch last-tile rescale: pipeline and flat (2)", + "Index-group tail and query state (1)", + "Alternating rescale and accumulate (1)", + "Query prefetch boundaries: 5, 15, 16 tiles (3)", + "Independent main/extra lengths with FP8/FP4 extra cache (4)" + ], + "schedule_comparisons": 8, + "baseline_regression": "ascending and nonadjacent source order failed before the fix", + "format": "all hooks passed" + }, + "schedule_cases": [ + { + "case": "bf16-2", + "old_II": 6057, + "new_II": 6059, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 416, + "new_strict_edges": 217, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "bf16-4", + "old_II": 6668, + "new_II": 6670, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 367, + "new_strict_edges": 192, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "bf16-10", + "old_II": 6687, + "new_II": 6689, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 351, + "new_strict_edges": 194, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "bf16-11", + "old_II": 6668, + "new_II": 6670, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 367, + "new_strict_edges": 198, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "bf16-16", + "old_II": 7514, + "new_II": 7516, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 391, + "new_strict_edges": 201, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "fp8", + "old_II": 7964, + "new_II": 7966, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 438, + "new_strict_edges": 226, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "fp4", + "old_II": 6185, + "new_II": 6187, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 438, + "new_strict_edges": 226, + "manual_order_violations": [], + "old_explicit_order_violations": [] + }, + { + "case": "flat", + "old_II": 8678, + "new_II": 8678, + "same_dependencies_and_storage": { + "data_deps": true, + "register_phase_deps": true, + "resource_deps": true, + "owner_exclusion_deps": true, + "pipe_order_deps": true, + "stage_order_deps": true, + "manual_stages": true, + "latencies": true, + "iis": true, + "resource_flags": true, + "buffer_sizes": true, + "memory_groups": true + }, + "manual_pipe_edges": 391, + "new_strict_edges": 0, + "manual_order_violations": [], + "old_explicit_order_violations": [] + } + ], + "gather_ss_schedule": [ + { + "id": 48, + "stage": 6, + "r": 4498, + "position": 46 + }, + { + "id": 49, + "stage": 10, + "r": 4499, + "position": 47 + }, + { + "id": 50, + "stage": 10, + "r": 4499, + "position": 48 + } + ], + "emitted_order": [ + { + "gather": 1219, + "delta_wait": 1226, + "ss_store": 1227 + }, + { + "gather": 1434, + "delta_wait": 1441, + "ss_store": 1442 + } + ], + "emitted_order_note": "Line numbers identify the generated AIV source retained with the local validation artifacts.", + "performance": { + "nq": 4096, + "nk": 4096, + "topk": 640, + "kv_format": "bf16", + "device": "Ascend950DT_9581", + "physical_device": 6, + "method": "kernelkit msprof, flush_l2=True, 3 alternating rounds of 10 launches", + "statistic": "Median of three round means; ten launches per round.", + "rounds": 3, + "launches_per_round": 10, + "round_mean_us": { + "before": [ + 866.1031723022461, + 866.2700653076172, + 865.6740188598633 + ], + "candidate": [ + 867.4383163452148, + 866.7469024658203, + 867.7959442138672 + ] + }, + "median_us": { + "before": 866.1031723022461, + "candidate": 867.4383163452148 + }, + "ratio_candidate_over_before": 1.0015415531147631, + "candidate_change_percent": 0.15415531147631167, + "bitwise_equal_to_baseline": true, + "source_sha256": { + "before": "b837367b6c3a0e9d8d9faabaac3946726361012ba6bf1b87050163a5b827425e", + "candidate": "986817d6f361116825ea0cb9a51db576e7a3fc9730c02936c1a0826fc4d826a2" + }, + "before_asc_sha256": "b6dcc15179098b2f42434d84a4b43b9faf019d8bcc7fb57c3d4b5b399a557a20", + "candidate_asc_sha256": "87e2df31c277e75d649622886ae55f88417733e982a34032231ae7ca23f2a284", + "baseline_library_sha256": "70a7dbb4d1e9eba8111985119c8d2960f0fc080c0604fbb2dcf8cbfc14858b80", + "baseline_description": "Intermediate reordered kernel retaining one IssuePhase; not the original e924481 or 9e691b0 source.", + "baseline_method": "Recompiled frozen old-compiler perf-final.asc via the Ascend postproc hook; main_kernel ABI and returned source identity checked.", + "isolation": "Timing began only after all 11 device regression tests exited successfully.", + "sample_limit": "The existing performance record saves round means, not individual launch samples." + }, + "coverage_limits": [ + "The latest record does not include a rerun of the full 14-case BF16 or nine-case decode correctness runners.", + "Timing covers one BF16 configuration; no updated three-case BF16 or ten-case decode performance sweep is recorded.", + "Earlier B4 and decode correctness/performance records retain their own source and compiler identities.", + "Modeled loop II is not measured device latency." + ], + "evidence_directory": ".tlcli/manual-pipe-order-20261020", + "evidence_files": [ + "summary.json", + "focused-tests.log", + "autoschedule-tests.log", + "device-tests.log", + "performance.json", + "performance.log", + "before.py", + "candidate.py", + "perf-before.asc", + "perf-candidate.asc" + ] +} diff --git a/examples/ascend/flashmla/run.py b/examples/ascend/flashmla/run.py new file mode 100644 index 0000000000..654e387bb9 --- /dev/null +++ b/examples/ascend/flashmla/run.py @@ -0,0 +1,195 @@ +"""Reproduce FlashMLA PR #229 with its input generator and msprof harness.""" + +from functools import partial +import argparse +import json +import os +from pathlib import Path +import statistics +import sys +import warnings + +warnings.filterwarnings("ignore", message="Permission mismatch.*") + +import torch +import torch_npu + + +def main(): + parser = argparse.ArgumentParser(__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--mode", choices=["smoke", "correctness", "bench"], default="smoke") + parser.add_argument("--implementation", choices=["both", "all", "reference", "manual", "staged"], default="both") + parser.add_argument("--output", type=Path, default=Path("results.json")) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--repeat", type=int, default=10) + parser.add_argument("--nk", type=int, help="Restrict the performance sweep to one KV length") + parser.add_argument("--detail", action="store_true", help="Also capture msprof pipe statistics") + parser.add_argument("--case", type=int, help="Run only this zero-based case number") + args = parser.parse_args() + sys.path[:0] = [str(args.reference.resolve()), str(args.reference.resolve() / "tests")] + import lib + import ref + import kernelkit as kk + from examples.ascend.flashmla import manual, staged + + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + torch.set_default_device("npu:0") + torch.set_default_dtype(torch.bfloat16) + torch.set_float32_matmul_precision("high") + if args.mode == "bench": + cases = [ + lib.TestParam(4096, nk, 640, h_q=64, have_attn_sink=True, seed=42) for nk in ([args.nk] if args.nk else [4096, 8192, 32768]) + ] + elif args.mode == "smoke": + cases = [lib.TestParam(32, 256, 128, h_q=64, have_attn_sink=True, seed=42)] + else: + cases = [ + lib.TestParam( + nq, + nk, + topk, + h_q=64, + seed=42, + have_attn_sink=sink, + have_topk_length=lengths, + k_amplifier_portion=0.02 if amplify else 0.0, + k_amplifier_ratio=256 if amplify else 1, + ) + for nq, nk, topk in [(1, 128, 128), (62, 95, 128), (213, 1521, 512), (65, 4096, 640)] + for sink, lengths, amplify in [(False, False, False), (True, True, False), (True, True, True)] + ] + cases += [ + lib.TestParam(65, 128, 64, h_q=64, seed=42, is_all_indices_invalid=True), + lib.TestParam(33, 95, 192, h_q=64, seed=42, have_attn_sink=True, have_topk_length=True), + ] + result = { + "torch": torch.__version__, + "torch_npu": torch_npu.__version__, + "device": torch.npu.get_device_name(0), + "visible_devices": os.environ.get("ASCEND_RT_VISIBLE_DEVICES"), + "cases": [], + } + args.output.parent.mkdir(parents=True, exist_ok=True) + for case_id, param in enumerate(cases): + if args.case is not None and case_id != args.case: + continue + print(f"CASE {param}", flush=True) + inputs = lib.generate_testcase(param) + if param.is_all_indices_invalid: + inputs.indices[::2].fill_(-1) + if args.mode == "correctness" and param.have_topk_length: + inputs.topk_length[0] = 0 + torch.npu.synchronize() + funcs = {} + if args.implementation in ("reference", "both", "all"): + funcs["reference"] = partial(lib.run_flash_mla_sparse_fwd, param, inputs) + for name, implementation in (("manual", manual), ("staged", staged)): + if args.implementation != name and args.implementation != "all" and not (args.implementation == "both" and name == "manual"): + continue + kernel = implementation.compile_prefill( + param.s_q, + param.s_kv, + param.topk, + sink=param.have_attn_sink, + variable_lengths=param.have_topk_length, + scale=inputs.sm_scale, + q_strides=tuple(inputs.q.stride()), + index_stride=inputs.indices.stride(0), + ) + source_path = args.output.parent / args.output.stem / f"{name}_{case_id}_{param.s_q}_{param.s_kv}_{param.topk}.asc" + source_path.parent.mkdir(parents=True, exist_ok=True) + source_path.write_text(kernel.get_kernel_source()) + print("Compiled cache:", getattr(kernel, "_tilelang_cache_path", None), flush=True) + sink_arg = inputs.attn_sink if inputs.attn_sink is not None else torch.zeros(64, dtype=torch.float32) + lengths_arg = inputs.topk_length if inputs.topk_length is not None else torch.full((param.s_q,), param.topk, dtype=torch.int32) + funcs[name] = partial(kernel, inputs.q, inputs.kv.view(param.s_kv, 512), inputs.indices.squeeze(1), sink_arg, lengths_arg) + outputs = {} + for name, fn in funcs.items(): + print(f"Launching {name}", flush=True) + outputs[name] = fn() + torch.npu.synchronize() + print(f"Finished {name}", flush=True) + case = { + "nq": param.s_q, + "nk": param.s_kv, + "topk": param.topk, + "sink": param.have_attn_sink, + "lengths": param.have_topk_length, + "amplified": param.k_amplifier_portion > 0, + "all_invalid_rows": param.is_all_indices_invalid, + "q_strides": list(inputs.q.stride()), + "index_strides": list(inputs.indices.stride()), + "correctness": {}, + } + if args.mode != "bench" or "reference" not in outputs: + expected_bf16, expected, expected_max, expected_lse = ref.ref_sparse_attn_fwd(param, inputs) + for name, (out, maximum, lse) in outputs.items(): + checks = [ + kk.check_is_allclose( + name + " out", + out.float(), + expected, + abs_tol=8e-4 if not param.k_amplifier_portion else 1.0, + rel_tol=3.01 / 128 if not param.k_amplifier_portion else 1.0, + cos_diff_tol=1e-5, + dtype_for_cos_diff_calc=torch.float, + ), + kk.check_is_allclose(name + " max", maximum, expected_max, abs_tol=1e-6, rel_tol=2.01 / 65536), + kk.check_is_allclose(name + " lse", lse, expected_lse, abs_tol=1e-5, rel_tol=4.01 / 65536), + ] + case["correctness"][name] = all(checks) + print(name, "max_abs_out", (out.float() - expected).abs().max().item(), flush=True) + if not all(case["correctness"].values()): + result["cases"].append(case) + args.output.write_text(json.dumps(result, indent=2)) + raise AssertionError("Correctness failure") + elif "reference" in outputs: + for name in outputs.keys() - {"reference"}: + for k in range(3): + torch.testing.assert_close( + outputs[name][k], outputs["reference"][k], rtol=0.03 if k == 0 else 1e-4, atol=8e-4 if k == 0 else 1e-5 + ) + case["correctness"][name + "_vs_reference"] = True + if args.mode == "bench": + flop = lib.count_flop_and_mem_vol(param, inputs).fwd_flop + samples = {name: [] for name in funcs} + launches = {name: [] for name in funcs} + for round_idx in range(args.rounds): + order = list(funcs) if round_idx % 2 == 0 else list(reversed(funcs)) + for name in order: + profile = kk.bench(funcs[name], num_tests=args.repeat, flush_l2=True) + needle = "sparse_attn_fwd" if name == "reference" else "main_kernel" + latency = profile.get_kernel_time(needle) + matching = [ranges for key, ranges in profile.time_ranges.items() if needle in key] + count = sum(len(ranges) for ranges in matching) + assert count == args.repeat, (name, count, profile.time_ranges.keys()) + samples[name].append(latency * 1e6) + launches[name].append([(end - start) * 1e6 for ranges in matching for start, end in ranges]) + print(f"{name} round {round_idx}: {latency * 1e6:.3f} us, {flop / latency / 1e12:.2f} TFLOPS", flush=True) + case["performance"] = { + name: { + "round_mean_us": values, + "launch_us": launches[name], + "median_us": statistics.median(values), + "tflops": flop / statistics.median(values) / 1e6, + } + for name, values in samples.items() + } + if args.detail: + from dataclasses import asdict + from tilelang.profiler import do_bench + + case["pipe_profile"] = {} + for name, fn in funcs.items(): + profile = do_bench(fn, backend="msprof_detail", early_stop_baseline=None, _n_warmup=5, _n_repeat=10) + case["pipe_profile"][name] = asdict(profile) + print(name, profile, flush=True) + result["cases"].append(case) + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/ascend/flashmla/run_decode.py b/examples/ascend/flashmla/run_decode.py new file mode 100644 index 0000000000..b1f2b88d30 --- /dev/null +++ b/examples/ascend/flashmla/run_decode.py @@ -0,0 +1,198 @@ +"""Same-input comparison against FlashMLA PR #229, including packed KV caches.""" + +from functools import partial +import argparse +import json +import os +from pathlib import Path +import statistics +import sys +import warnings + +warnings.filterwarnings("ignore", message="Permission mismatch.*") +import torch +import torch_npu + + +def main(): + parser = argparse.ArgumentParser(__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--mode", choices=["smoke", "correctness", "bench"], default="smoke") + parser.add_argument("--implementation", choices=["both", "all", "reference", "manual", "staged"], default="both") + parser.add_argument("--output", type=Path, default=Path("decode-results.json")) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--repeat", type=int, default=10) + parser.add_argument("--case", type=int, help="Run only this zero-based case number") + args = parser.parse_args() + sys.path[:0] = [str(args.reference.resolve()), str(args.reference.resolve() / "tests")] + import lib + import ref + import quant + import flash_mla + import kernelkit as kk + from examples.ascend.flashmla.decode import prepare_decode + + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + torch.set_default_device("npu:0") + torch.set_default_dtype(torch.bfloat16) + torch.set_float32_matmul_precision("high") + fp8 = quant.KVCacheLayout.V41_FP8Sparse + fp4 = quant.KVCacheLayout.V41_FP4Sparse + raw = lib.RawTestParamForDecode + if args.mode == "bench": + cases = [ + raw( + b, + 64, + 4, + 1, + 256, + True, + 128, + extra_s_k=2048, + extra_topk=512, + block_size=256, + extra_block_size=64, + kvcache_layout=fp8, + extra_kvcache_layout=fmt, + seed=42, + ) + for fmt in [fp8, fp4] + for b in [64, 128, 256, 512] + ] + cases += [ + raw( + 1, + 64, + 4096, + 1, + 4096, + False, + 128, + extra_s_k=4096, + extra_topk=512, + block_size=64, + extra_block_size=64, + kvcache_layout=fp8, + extra_kvcache_layout=fmt, + seed=42, + ) + for fmt in [fp8, fp4] + ] + else: + cases = [raw(4, 64, 1, 1, 256, True, 128, kvcache_layout=fp8, seed=42)] + cases += [ + raw(4, 64, 3, 1, 256, True, 128, extra_s_k=512, extra_topk=128, kvcache_layout=fp8, extra_kvcache_layout=fmt, seed=42) + for fmt in [fp8, fp4] + ] + if args.mode == "correctness": + cases += [ + raw( + b, + 64, + 3, + 1, + 650, + True, + 576, + extra_s_k=512, + extra_topk=64, + block_size=53, + extra_block_size=61, + have_topk_length=lengths, + have_extra_topk_length=lengths, + have_zero_seqlen_k=True, + enable_attn_sink=sink, + is_all_indices_invalid=invalid, + kvcache_layout=fp8, + extra_kvcache_layout=fmt, + seed=42, + ) + for fmt in [fp8, fp4] + for b, lengths, sink, invalid in [(1, False, False, True), (74, True, True, False), (35, False, False, False)] + ] + result = { + "torch": torch.__version__, + "torch_npu": torch_npu.__version__, + "device": torch.npu.get_device_name(0), + "visible_devices": os.environ.get("ASCEND_RT_VISIBLE_DEVICES"), + "cases": [], + } + args.output.parent.mkdir(parents=True, exist_ok=True) + for case_id, rp in enumerate(cases): + if args.case is not None and case_id != args.case: + continue + p = rp.to_test_param() + print(f"CASE {case_id}: {p}", flush=True) + inputs = lib.generate_testcase_for_decode(p) + torch.npu.synchronize() + metadata, _ = flash_mla.get_mla_metadata() + funcs = {} + if args.implementation in ("reference", "both", "all"): + funcs["reference"] = partial(lib.run_flash_mla_decode, p, inputs, metadata, None) + for name in ("manual", "staged"): + if args.implementation != name and args.implementation != "all" and not (args.implementation == "both" and name == "manual"): + continue + kernel, funcs[name] = prepare_decode(p, inputs, implementation=name) + source_path = args.output.parent / args.output.stem / f"{name}_decode_{case_id}.asc" + source_path.parent.mkdir(parents=True, exist_ok=True) + source_path.write_text(kernel.get_kernel_source()) + print("Compiled cache:", getattr(kernel, "_tilelang_cache_path", None), flush=True) + outputs = {name: fn() for name, fn in funcs.items()} + torch.npu.synchronize() + expected = outputs["reference"] if args.mode == "bench" and "reference" in outputs else ref.ref_sparse_attn_decode(p, inputs) + case = { + "id": case_id, + "b": rp.b, + "sq": rp.s_q, + "topk": rp.topk, + "extra_topk": rp.extra_topk, + "extra_format": "fp4" if rp.extra_kvcache_layout == fp4 else "fp8", + "correctness": {}, + } + for name, (out, lse) in outputs.items(): + checks = [ + kk.check_is_allclose(name + " out", out, expected[0], abs_tol=1e-3, rel_tol=2.01 / 128, cos_diff_tol=5e-6), + kk.check_is_allclose(name + " lse", lse, expected[1], abs_tol=1e-6, rel_tol=8.01 / 65536), + ] + case["correctness"][name] = all(checks) + print( + name, + "out error", + (out - expected[0]).abs().max().item(), + "lse error", + (lse - expected[1]).abs().nan_to_num().max().item(), + flush=True, + ) + assert all(case["correctness"].values()), case + if args.mode == "bench": + flop = lib.count_flop_and_mem_vol_for_decode(p, inputs).flop + samples = {name: [] for name in funcs} + launches = {name: [] for name in funcs} + for round_idx in range(args.rounds): + for name in list(funcs) if round_idx % 2 == 0 else list(reversed(funcs)): + profile = kk.bench(funcs[name], num_tests=args.repeat, flush_l2=True) + needle = "sparse_attn_fwd" if name == "reference" else "main_kernel" + latency = profile.get_kernel_time(needle) + matching = [ranges for key, ranges in profile.time_ranges.items() if needle in key] + assert sum(len(ranges) for ranges in matching) == args.repeat + samples[name].append(latency * 1e6) + launches[name].append([(end - start) * 1e6 for ranges in matching for start, end in ranges]) + print(f"{name} round {round_idx}: {latency * 1e6:.3f} us, {flop / latency / 1e12:.2f} TFLOPS", flush=True) + case["performance"] = { + name: { + "round_mean_us": values, + "launch_us": launches[name], + "median_us": statistics.median(values), + "tflops": flop / statistics.median(values) / 1e6, + } + for name, values in samples.items() + } + result["cases"].append(case) + args.output.write_text(json.dumps(result, indent=2)) + print("PASS", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/ascend/flashmla/staged.py b/examples/ascend/flashmla/staged.py new file mode 100644 index 0000000000..ee9e49b7c4 --- /dev/null +++ b/examples/ascend/flashmla/staged.py @@ -0,0 +1,757 @@ +"""FlashMLA with logical tiles and compiler-managed buffer versions. + +A persistent KV stream spans query jobs. Q and the online state live until +the query's last KV tile; output chunks own their Fixpipe/Vector handoff. +Q uses two explicit query slots when its periodic ownership fits the event +budget; all per-tile buffer versions are compiler-managed. +""" + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.layout import make_ascend_compact_nz_layout, make_ascend_l0c_layout, make_ascend_major_k_layout, make_ascend_major_mn_layout +from examples.ascend.flashmla.quantization import dequantize + +PASS_CONFIG = { + tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: True, + tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, +} + + +def rolled(n): + return T.unroll(n, annotations={"pragma_unroll_factor": 1}) + + +def sparse_attention( + nq, + nk, + topk, + *, + sink=True, + variable_lengths=False, + scale=512 ** (-0.5), + q_strides=None, + index_stride=None, + kv_format="bf16", + extra_topk=0, + extra_format="fp8", + extra_index_stride=None, + variable_extra_lengths=False, + kv_storage_bytes=1, + extra_storage_bytes=1, + cache_hint=0, + pipeline_stages=True, +): + if nq <= 0 or nk <= 0 or topk <= 0 or topk % 64: + raise ValueError("nq/nk must be positive and topk a positive multiple of 64") + if kv_format not in ("bf16", "fp8") or extra_format not in ("fp8", "fp4"): + raise ValueError("Supported formats are BF16 prefill and FP8 main / FP8 or FP4 extra decode") + if extra_topk < 0 or extra_topk % 64: + raise ValueError("extra_topk must be a nonnegative multiple of 64") + if max(kv_storage_bytes, extra_storage_bytes) >= 2**32: + raise ValueError("KV byte offsets must fit uint32") + decode = kv_format != "bf16" + extra_fp4 = extra_format == "fp4" + valid_limit = T.uint32(2147483648) if decode else nk + q_strides = q_strides or (64 * 512, 512, 1) + index_stride = index_stride or topk + extra_index_stride = extra_index_stride or max(extra_topk, 1) + nb0 = topk // 64 + nb = nb0 + extra_topk // 64 + # Explicit query slots repeat every 2*nb lexical KV iterations. + # The two AIC Q ownership directions need 2*(nb+1) mutex IDs, capped at 32. + q_prefetch = pipeline_stages and 4 <= nb <= 15 + q_slots = 2 if q_prefetch else 1 + q_lead_shift = nb - 4 if q_prefetch else 0 + extra_metadata_lead = pipeline_stages and q_lead_shift > 0 + index_group = min(640, max(topk, extra_topk)) + gather_cache_policy = "normal_lv" if not decode or cache_hint else "notalloc_clean" + direct_query_masks = not decode and pipeline_stages and extra_topk == 0 and 5 <= nb0 <= 10 + query_mask_stride = index_group // 2 + 32 + + @T.macro + def pv_gemm(k, unit_flag, clear, d, aq, bv, cpv): + # Keep the zero-flag path explicit so intermediate rescale drains + # retain their ordinary M/FIX synchronization. + with T.Task(): + if (unit_flag != 0) if nb == 2 else (k == 0 or k == nb - 1): + T.gemm( + aq[:, :64], + bv, + cpv[d, :, :], + clear_accum=clear, + unit_flag_ctrl=unit_flag if nb == 2 else T.if_then_else(k == nb - 1, 3, 2), + ) + else: + T.gemm(aq[:, :64], bv, cpv[d, :, :], clear_accum=clear, unit_flag_ctrl=0) + + @T.macro + def softmax_vf(first, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta, mask_base=0): + with T.SimdVF(): + full = T.simd.pset(32) + eight = T.simd.pset(32, "PAT_VL8") + one = T.simd.pset(32, "PAT_VL1") + masks = T.alloc_local((4,), "boolx256") + for col in T.unroll(4): + lo = T.simd.vld(T.access_ptr(smaller[mask_base + col * 8], "r", extent=8), "BLK") + hi = T.simd.vld(T.access_ptr(larger[mask_base + col * 8], "r", extent=8), "BLK") + paired, ignored = T.simd.vintlv(lo, hi) + masks[col] = T.simd.vcmps(paired, valid_limit, op="lt") + old_sum = T.simd.vld(sums[0]) + old_scale_max = T.simd.vld(scaling_max[0]) + for r in rolled(8): + rm = T.alloc_var("float32x64") + rm = T.simd.vdup(-3.4028234663852886e38, "float32") + if not first: + rm = T.simd.vld(maxima[r * 8], "E2B_B32") + for col in T.unroll(4): + x = T.simd.vld(p[r * 4, col * 16]) + rm = T.simd.vmax(rm, x, masks[col], mode="MODE_MERGING") + reduced = T.simd.vcgmax(rm) + even, odd = T.simd.vdintlv(reduced, reduced) + rmax = T.simd.vmax(even, odd) + repeated, ignored2 = T.simd.vintlv(rmax, rmax) + T.simd.vsts(maxima[r * 8], repeated, eight) + T.simd.mem_bar("VST_VLD") + new_max = T.simd.vld(maxima[0]) + chosen_max = T.alloc_var("float32x64") + factor = T.alloc_var("float32x64") + chosen_max = new_max + factor = T.simd.vdup(0.0, "float32") + if not first: + diff = T.simd.vsub(new_max, old_scale_max) + greatest = T.simd.vcmax(diff) + broadcast = T.simd.vdupv(greatest) + rescale = T.simd.vcmps(broadcast, 6.0 / scale, op="gt") + chosen_max = T.simd.vsel(new_max, old_scale_max, rescale) + factor = T.simd.vexpdif(T.simd.vmuls(old_scale_max, scale), T.simd.vmuls(chosen_max, scale)) + T.simd.vsts(delta[0], greatest, one) + T.simd.vsts(factors[0], factor) + T.simd.vsts(scaling_max[0], chosen_max) + T.simd.mem_bar("VST_VLD") + for r in rolled(4): + rmax2 = T.alloc_local((2,), "float32x64") + rsum2 = T.alloc_local((2,), "float32x64") + packed2 = T.alloc_local((2, 2), "bfloat16x128") + for rr in T.unroll(2): + rmax2[rr] = T.simd.vmuls(T.simd.vld(scaling_max[r * 16 + rr * 8], "E2B_B32"), scale) + for pair in T.unroll(2): + for cc in T.unroll(2): + for rr in T.unroll(2): + x = T.simd.vld(p[r * 8 + rr * 4, (pair * 2 + cc) * 16]) + prob = T.simd.vexpdif(T.simd.vmuls(x, scale), rmax2[rr], masks[pair * 2 + cc]) + if pair == 0 and cc == 0: + rsum2[rr] = prob + else: + rsum2[rr] = T.simd.vadd(rsum2[rr], prob) + if cc == 0: + packed2[rr, cc] = T.simd.vcvt(prob, "bfloat16", full, part=0) + else: + packed2[rr, cc] = T.simd.vcvt(prob, "bfloat16", full, part=1) + merged0 = T.simd.vadd(packed2[0, 0], packed2[0, 1]) + merged1 = T.simd.vadd(packed2[1, 0], packed2[1, 1]) + col0, col1 = T.simd.vdintlv(merged0, merged1) + T.simd.vsts(s[r * 8, pair * 32], col0) + T.simd.vsts(s[r * 8, pair * 32 + 16], col1) + for rr in T.unroll(2): + reduced_sum = T.simd.vcpadd(T.simd.vcgadd(rsum2[rr])) + repeated_sum, ignored3 = T.simd.vintlv(reduced_sum, reduced_sum) + T.simd.vsts(sums[r * 16 + rr * 8], repeated_sum, eight) + T.simd.mem_bar("VST_VLD") + if not first: + current_sum = T.simd.vld(sums[0]) + updated_sum = T.simd.vadd(T.simd.vmul(old_sum, factor), current_sum) + T.simd.vsts(sums[0], updated_sum) + + @T.macro + def scale_output_vf(add, d, factors, out, accum): + with T.SimdVF(): + for r in rolled(32): + factor5 = T.simd.vld(factors[r * 2], "BRC_B32") + for j in T.unroll(2): + cur = T.alloc_var("float32x64", init=T.simd.vld(out[r, j * 64])) + if add: + old = T.simd.vld(accum[r, d * 128 + j * 64]) + cur = T.simd.vadd(old, cur) + T.simd.vsts(accum[r, d * 128 + j * 64], T.simd.vmul(cur, factor5)) + + @T.macro + def final_output_vf(add, d, denorm, out, accum, final): + with T.SimdVF(): + mask16 = T.simd.pset(16) + mask32 = T.simd.pset(32) + for r in rolled(16): + packed = T.alloc_local((2, 2), "bfloat16x128") + for rr in T.unroll(2): + norm = T.simd.vld(denorm[r * 2 + rr], "BRC_B32") + for cc in T.unroll(2): + x = T.simd.vld(out[r * 2 + rr, cc * 64]) + if add: + old = T.simd.vld(accum[r * 2 + rr, d * 128 + cc * 64]) + x = T.simd.vadd(x, old, mask32) + scaled = T.simd.vmul(x, norm, mask32) + if rr == 0: + packed[rr, cc] = T.simd.vcvt(scaled, "bfloat16", mask32, part=0) + else: + packed[rr, cc] = T.simd.vcvt(scaled, "bfloat16", mask32, part=1) + merged0 = T.simd.vadd(packed[0, 0], packed[1, 0], mask16) + merged1 = T.simd.vadd(packed[0, 1], packed[1, 1], mask16) + row0, row1 = T.simd.vdintlv(merged0, merged1) + T.simd.vsts(final[r * 2, d * 128], row0, mask16) + T.simd.vsts(final[r * 2 + 1, d * 128], row1, mask16) + + @T.macro + def load_metadata( + k, job, qidx, extra, Indices, ExtraIndices, Lengths, ExtraLengths, indices, group_smaller, group_larger, smaller, larger + ): + rel = T.alloc_var("int32") + rel = T.if_then_else(extra, k - nb0, k) + group_tile = T.alloc_var("int32") + group_tile = rel % (index_group // 64) + if k == 0 if direct_query_masks else group_tile == 0: + if extra_topk > 0 and extra: + group_length = T.min(index_group, extra_topk - rel * 64) + T.copy(ExtraIndices[qidx, rel * 64 : rel * 64 + group_length], indices[:group_length]) + else: + group_length = T.min(index_group, topk - rel * 64) + T.copy(Indices[qidx, rel * 64 : rel * 64 + group_length], indices[:group_length]) + main_length = T.min(topk, Lengths[qidx]) if variable_lengths else topk + extra_length = T.min(extra_topk, ExtraLengths[qidx]) if variable_extra_lengths else extra_topk + valid_len = T.max(T.if_then_else(extra, extra_length, main_length) - rel * 64, 0) + with T.SimdVF(): + pairs32 = T.simd.pset(32, "PAT_VL32") + for group_k in rolled(index_group // 64): + x = T.reinterpret("uint32x64", T.simd.vld(indices[group_k * 64])) + pos = T.simd.vci(group_k * 64, "int32") + x = T.simd.vsel(x, T.simd.vdup(valid_limit, "uint32"), T.simd.vcmps(pos, valid_len, op="lt")) + even, odd = T.simd.vdintlv(x, x) + lo = T.simd.vmins(T.simd.vmin(even, odd), valid_limit) + hi = T.simd.vmax(even, odd) + if direct_query_masks: + T.simd.vsts(group_smaller[(job % 2) * query_mask_stride + group_k * 32], lo, pairs32, extent=32) + T.simd.vsts(group_larger[(job % 2) * query_mask_stride + group_k * 32], hi, pairs32, extent=32) + else: + T.simd.vsts(group_smaller[group_k * 32], lo, pairs32, extent=32) + T.simd.vsts(group_larger[group_k * 32], hi, pairs32, extent=32) + if not direct_query_masks: + # Snapshot only the active tile's masks; the index group may + # be replaced while a previous tile is still in softmax. + with T.SimdVF(): + pairs32 = T.simd.pset(32, "PAT_VL32") + T.simd.vsts(smaller[0], T.simd.vld(group_smaller[group_tile * 32]), pairs32, extent=32) + T.simd.vsts(larger[0], T.simd.vld(group_larger[group_tile * 32]), pairs32, extent=32) + + @T.macro + def gather_kv(k, job, sid, extra, KV, ExtraKV, group_smaller, group_larger, smaller, larger, kv, raw8, raw4): + for i in T.unroll(16): + with T.Stage(0): + if direct_query_masks: + lo = group_smaller[(job % 2) * query_mask_stride + k * 32 + sid * 16 + i] + hi = group_larger[(job % 2) * query_mask_stride + k * 32 + sid * 16 + i] + else: + lo = smaller[sid * 16 + i] + hi = larger[sid * 16 + i] + second = T.if_then_else(hi < valid_limit, hi, lo) + n_burst = T.alloc_var("int32") + n_burst = T.if_then_else(lo < valid_limit, 2, 0) + safe_lo = T.if_then_else(n_burst == 0, T.uint32(0), lo) + stride = T.alloc_var("int64") + stride = T.Cast("int64", second - lo) + if decode and extra_topk > 0 and extra_fp4: + # Both formats retain separate bytes, sharing one + # MTE2 task and one V consumer ownership boundary. + with T.Task(): + if extra: + extra_pairs = T.StridedTensor( + (extra_storage_bytes // 288, 2, 288), + (288, stride * 288, 1), + "uint8", + data=ExtraKV.data, + ) + for _ in T.Serial(1): + T.assume(n_burst > 0) + T.assume(n_burst == 0 or safe_lo < extra_pairs.shape[0]) + T.copy( + extra_pairs[safe_lo, :n_burst, :], + raw4[i * 2 : i * 2 + n_burst, :], + l2_cache_ctrl=gather_cache_policy, + ) + else: + main_pairs = T.StridedTensor( + (kv_storage_bytes // 528, 2, 528), + (528, stride * 528, 1), + "uint8", + data=KV.data, + ) + for _ in T.Serial(1): + T.assume(n_burst > 0) + T.assume(n_burst == 0 or safe_lo < main_pairs.shape[0]) + T.copy( + main_pairs[safe_lo, :n_burst, :], + raw8[i * 2 : i * 2 + n_burst, :528], + l2_cache_ctrl=gather_cache_policy, + ) + elif decode: + if extra_topk > 0 and extra: + extra_pairs = T.StridedTensor( + (extra_storage_bytes // 528, 2, 528), + (528, stride * 528, 1), + "uint8", + data=ExtraKV.data, + ) + # Limit DMA assumptions to this copy. + for _ in T.Serial(1): + with T.Stage(0): + T.assume(n_burst > 0) + T.assume(n_burst == 0 or safe_lo < extra_pairs.shape[0]) + T.copy( + extra_pairs[safe_lo, :n_burst, :], + raw8[i * 2 : i * 2 + n_burst, :528], + l2_cache_ctrl=gather_cache_policy, + ) + else: + main_pairs = T.StridedTensor( + (kv_storage_bytes // 528, 2, 528), + (528, stride * 528, 1), + "uint8", + data=KV.data, + ) + # Limit DMA assumptions to this copy. + for _ in T.Serial(1): + with T.Stage(0): + T.assume(n_burst > 0) + T.assume(n_burst == 0 or safe_lo < main_pairs.shape[0]) + T.copy( + main_pairs[safe_lo, :n_burst, :], + raw8[i * 2 : i * 2 + n_burst, :528], + l2_cache_ctrl=gather_cache_policy, + ) + else: + main_pairs = T.StridedTensor( + (nk, 2, 512), + (512, stride * 512, 1), + "bfloat16", + data=KV.data, + ) + # Limit DMA assumptions to this copy. + for _ in T.Serial(1): + with T.Stage(0): + T.assume(n_burst > 0) + T.assume(n_burst == 0 or safe_lo < main_pairs.shape[0]) + T.copy( + main_pairs[safe_lo, :n_burst, :], + kv[i * 2 : i * 2 + n_burst, :], + l2_cache_ctrl=gather_cache_policy, + ) + + @T.macro + def attention_body(Q, KV, Indices, Sink, Lengths, ExtraKV, ExtraIndices, ExtraLengths, O, Max, LSE): + with T.MixedKernel(min(nq, 32)) as (bx, sid): + q_l1 = T.alloc_l1((q_slots, 64, 512), "bfloat16") + kv_l1_0 = T.alloc_l1((32, 512), "bfloat16") + kv_l1_1 = T.alloc_l1((32, 512), "bfloat16") + s_l1 = T.alloc_l1((64, 64), "bfloat16") + aq = T.alloc_l0a((64, 128), "bfloat16") + bq = T.alloc_l0b((64, 128), "bfloat16") + bv = T.alloc_l0b((64, 128), "bfloat16") + cq = T.alloc_l0c((64, 64), "float32") + # Four disjoint output-column tiles, each accumulated across all KV tiles. + cpv = T.alloc_l0c((4, 64, 128), "float32") + kv = T.alloc_shared((32, 1 if decode else 512), "bfloat16") + raw8 = T.alloc_shared((32, 544 if decode else 1), "uint8") + raw4 = T.alloc_shared((32, 288 if decode and extra_fp4 and extra_topk else 1), "uint8") + scales = T.alloc_shared((1024 if decode else 1,), "bfloat16") + nz = T.alloc_shared((33, 512), "bfloat16") + p = T.alloc_shared((32, 64), "float32") + s = T.alloc_shared((32, 64), "bfloat16") + accum = T.alloc_shared((32, 512), "float32") + out = T.alloc_shared((32, 128), "float32") + # Pack consumed FP32 chunks into the row's first 512 BF16 elements. + # Drain the complete row only after every d chunk is normalized. + final = T.view(accum, shape=(32, 1024), dtype="bfloat16") + maxima = T.alloc_shared((64,), "float32") + scaling_max = T.alloc_shared((64,), "float32") + sums = T.alloc_shared((64,), "float32") + denorm = T.alloc_shared((64,), "float32") + factors = T.alloc_shared((64,), "float32") + delta = T.alloc_shared((64,), "float32") + attn_sink = T.alloc_shared((64,), "float32") + indices = T.alloc_shared((index_group,), "int32") + group_smaller = T.alloc_shared(((2 if direct_query_masks else 1) * query_mask_stride,), "uint32") + group_larger = T.alloc_shared(((2 if direct_query_masks else 1) * query_mask_stride,), "uint32") + if not direct_query_masks: + smaller = T.alloc_shared((32,), "uint32") + larger = T.alloc_shared((32,), "uint32") + ss = T.alloc_ss((2,), "uint32") + ss_pairs = T.view(ss, shape=(1,), dtype="uint64") + ss_result = T.alloc_ss((1,), "uint32") + T.annotate_l1_offsets( + { + kv_l1_0: 0, + q_l1: 131072, + kv_l1_1: 262144, + s_l1: 393216, + } + ) + T.annotate_layout( + { + aq: make_ascend_major_k_layout(aq), + bq: make_ascend_major_k_layout(bq), + bv: make_ascend_major_mn_layout(bv), + q_l1: make_ascend_major_k_layout(q_l1), + kv_l1_0: make_ascend_major_k_layout(kv_l1_0), + kv_l1_1: make_ascend_major_k_layout(kv_l1_1), + s_l1: make_ascend_major_k_layout(s_l1), + nz: make_ascend_compact_nz_layout(nz), + s: make_ascend_compact_nz_layout(s), + p: make_ascend_l0c_layout(p), + } + ) + if q_prefetch: + T.annotate_manual_multi_buffer({q_l1: 2 * nb}) + else: + T.annotate_buffer_versions({q_l1: 1}) + T.annotate_buffer_versions( + { + kv_l1_0: 4, + kv_l1_1: 4, + s_l1: 2, + aq: 4, + bq: 4, + bv: 4, + cq: 1, + cpv: 1, + kv: 1 if decode else 3, + raw8: 3 if decode else 1, + raw4: 3 if decode and extra_fp4 and extra_topk else 1, + nz: 1, + scales: 1, + p: 1, + s: 1, + accum: 1, + out: 2, + final: 1, + maxima: 1, + scaling_max: 1, + sums: 1, + denorm: 2, + factors: 2, + delta: 1, + indices: 1, + group_smaller: 1, + group_larger: 1, + **( + {} + if direct_query_masks + else { + smaller: 6 if extra_metadata_lead else 5, + larger: 6 if extra_metadata_lead else 5, + } + ), + attn_sink: 1, + ss: 2, + ss_result: 1, + } + ) + if direct_query_masks: + T.annotate_manual_multi_buffer({group_smaller: 2 * nb, group_larger: 2 * nb}) + if sink: + T.copy(Sink[sid * 32 : sid * 32 + 32], attn_sink[:32]) + # Initialize every automatic version once; zero-burst DMA preserves it. + with T.SimdVF(): + if decode: + T.fill(raw8, 0) + if extra_topk > 0 and extra_fp4: + T.fill(raw4, 0) + else: + T.fill(kv, 0) + emitted = T.alloc_var("int32", init=0) + # Each nested scheduler loop is manual independently. VF loops + # belong to their enclosing VF task and cannot contain T.Stage. + for step in T.Pipelined(T.ceildiv(nq - bx, 32) * nb, num_stages=6 + q_lead_shift, annotations={"enable_offset": True}): + T.assume_conflict(bq, bv) + job = step // nb + k = step % nb + qidx = bx + job * 32 + with T.Stage(0): + if q_prefetch and k == 0: + T.copy(Q[qidx, :, :], q_l1[job % q_slots, :, :]) + # Flat execution needs the producer tiles before their consumers. + if not pipeline_stages: + with T.Stage(q_lead_shift - (1 if extra_metadata_lead else 0)): + extra = k >= nb0 + load_metadata( + k, + job, + qidx, + extra, + Indices, + ExtraIndices, + Lengths, + ExtraLengths, + indices, + group_smaller, + group_larger, + None if direct_query_masks else smaller, + None if direct_query_masks else larger, + ) + with T.Stage(q_lead_shift + (0 if extra_metadata_lead else 1) if pipeline_stages else 0): + gather_kv( + k, + job, + sid, + extra, + KV, + ExtraKV, + group_smaller, + group_larger, + None if direct_query_masks else smaller, + None if direct_query_masks else larger, + kv, + raw8, + raw4, + ) + # In the pipeline, source order follows issue order; Stage still + # identifies the logical tile used by dependency analysis. + with T.Stage(2 + q_lead_shift if pipeline_stages else 0): + if decode: + if extra_topk > 0 and extra_fp4: + with T.Task(): + if k >= nb0: + dequantize(raw4, scales, nz, None, True) + else: + dequantize(raw8, scales, nz, None, False) + else: + dequantize(raw8, scales, nz, None, False) + else: + with T.SimdVF(): + nd_mask = T.simd.pset(16) + src_ptr = T.simd.make_ubuf_ptr(T.access_ptr(kv[0, 0], "r", extent=16384), "bfloat16") + dst0 = T.simd.make_ubuf_ptr(T.access_ptr(nz[0, 0], "w", extent=16896), "bfloat16") + dst1 = T.simd.make_ubuf_ptr(T.access_ptr(nz[0, 128], "w", extent=12672), "bfloat16") + for _col in rolled(2): + for row in T.unroll(32): + x0, src_ptr = T.simd.vld(src_ptr, post_inc=128) + x1, src_ptr = T.simd.vld(src_ptr, post_inc=T.if_then_else(row == 31, -15744, 384)) + config = (33 << 16) | T.if_then_else(row == 31, 497, 1) + dst0 = T.simd.vsstb(x0, dst0, config, nd_mask, update=True) + dst1 = T.simd.vsstb(x1, dst1, config, nd_mask, update=True) + if sid == 0: + T.copy(nz[:32, :], kv_l1_0) + else: + T.copy(nz[:32, :], kv_l1_1) + with T.Stage(4 + q_lead_shift if pipeline_stages else 0): + # Adjacent Cube configuration regions follow physical phase. + if not q_prefetch and k == 0: + T.copy(Q[qidx, :, :], q_l1[0, :, :]) + T.set_mmad_direction("m") + for d in T.Pipelined(4, num_stages=4): + with T.Stage(0): + T.copy(q_l1[job % q_slots, :, d * 128 : (d + 1) * 128], aq) + with T.Task(): + T.copy(kv_l1_0[:, d * 128 : (d + 1) * 128], bq[:32, :], l0_pitch=4) + T.copy(kv_l1_1[:, d * 128 : (d + 1) * 128], bq[32:, :], l0_pitch=4) + T.gemm(aq, bq, cq, transpose_B=True, clear_accum=d == 0, unit_flag_ctrl=T.if_then_else(d == 3, 3, 2)) + T.dual_copy(cq, p, unit_flag_ctrl=3) + with T.Stage(4 + q_lead_shift if pipeline_stages else 0): + if k == 0: + with T.SimdVF(): + T.fill(maxima, -3.4028234663852886e38) + T.fill(scaling_max, 0.0) + T.fill(sums, 0.0) + delta[0] = T.float32(0) + if k == 0: + if direct_query_masks: + softmax_vf( + True, + group_smaller, + group_larger, + p, + s, + maxima, + sums, + scaling_max, + factors, + delta, + (job % 2) * query_mask_stride + k * 32, + ) + else: + softmax_vf(True, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta) + else: + if direct_query_masks: + softmax_vf( + False, + group_smaller, + group_larger, + p, + s, + maxima, + sums, + scaling_max, + factors, + delta, + (job % 2) * query_mask_stride + k * 32, + ) + else: + softmax_vf(False, smaller, larger, p, s, maxima, sums, scaling_max, factors, delta) + T.copy(s[:32, :], s_l1[sid * 32 : sid * 32 + 32, :]) + if not pipeline_stages: + ss[sid] = T.Cast("uint32", delta[0] > 6.0 / scale) + with T.Stage(4 + q_lead_shift if pipeline_stages else 0): + if k == nb - 1: + with T.SimdVF(): + rmax, unused = T.simd.vdintlv(T.simd.vld(maxima[0]), T.simd.vld(maxima[0])) + smax, unused1 = T.simd.vdintlv(T.simd.vld(scaling_max[0]), T.simd.vld(scaling_max[0])) + rsum, unused2 = T.simd.vdintlv(T.simd.vld(sums[0]), T.simd.vld(sums[0])) + invalid = T.simd.vcmps(rmax, -3.4028234663852886e38, op="eq") + maximum = T.simd.vmuls(rmax, scale) + scaling = T.simd.vmuls(smax, scale) + logsum = T.simd.vadd(T.simd.vln(rsum), scaling) + denom = T.alloc_var("float32x64", init=rsum) + if sink: + sinkprob = T.simd.vexpdif(T.simd.vld(attn_sink[0]), scaling) + denom = T.simd.vadd(denom, sinkprob) + reciprocal = T.simd.vdiv(T.simd.vdup(1.0, "float32"), denom) + T.simd.vsts(denorm[0], T.simd.vsel(T.simd.vdup(0.0, "float32"), reciprocal, invalid)) + T.simd.vsts(maxima[0], T.simd.vsel(T.simd.vdup(float("-inf"), "float32"), maximum, invalid)) + T.simd.vsts(sums[0], T.simd.vsel(T.simd.vdup(float("inf"), "float32"), logsum, invalid)) + if not decode: + T.copy(maxima[:32], Max[qidx, sid * 32 : sid * 32 + 32]) + T.copy(sums[:32], LSE[qidx, sid * 32 : sid * 32 + 32]) + if pipeline_stages: + with T.Stage(q_lead_shift - (1 if extra_metadata_lead else 0)): + extra = k >= nb0 + load_metadata( + k, + job, + qidx, + extra, + Indices, + ExtraIndices, + Lengths, + ExtraLengths, + indices, + group_smaller, + group_larger, + None if direct_query_masks else smaller, + None if direct_query_masks else larger, + ) + with T.Stage(5 + q_lead_shift if pipeline_stages else 0): + if k == 0: + emitted = 0 + T.set_mmad_direction("n") + # Bind this pure predicate before the PV owner also on PIPE_S. + last_output = k == nb - 1 + if nb == 2: + # Keep the zero-flag alternative visible to sync + # analysis when both KV tiles are boundaries. + pv_flag = T.alloc_var("int32", init=T.if_then_else(k == nb - 1, 3, 2)) + for _pv_region in T.Pipelined(1, num_stages=4): + with T.Stage(0): + T.copy(s_l1, aq[:, :64]) + clear = T.alloc_var("int32", init=T.Cast("int32", ss_pairs[0] != T.uint64(0))) + ss_result[0] = T.Cast("uint32", clear) + for _pv_branches in T.serial(1): + if ss_result[0] != T.uint32(0): + for d in T.Pipelined(4, num_stages=2): + with T.Stage(0): + T.dual_copy(cpv[d, :, :], out) + if emitted > 0: + scale_output_vf(True, d, factors, out, accum) + else: + scale_output_vf(False, d, factors, out, accum) + emitted = emitted + 1 + for d in T.Pipelined(4, num_stages=2): + with T.Stage(0): + with T.Task(): + T.copy(kv_l1_0[:, d * 128 : (d + 1) * 128], bv[:32, :], l0_pitch=8) + T.copy(kv_l1_1[:, d * 128 : (d + 1) * 128], bv[32:, :], l0_pitch=8) + pv_gemm(k, pv_flag if nb == 2 else 0, True, d, aq, bv, cpv) + else: + for d in T.Pipelined(4, num_stages=2): + with T.Stage(0): + with T.Task(): + T.copy(kv_l1_0[:, d * 128 : (d + 1) * 128], bv[:32, :], l0_pitch=8) + T.copy(kv_l1_1[:, d * 128 : (d + 1) * 128], bv[32:, :], l0_pitch=8) + pv_gemm(k, pv_flag if nb == 2 else 0, k == 0, d, aq, bv, cpv) + with T.Stage(5 + q_lead_shift if pipeline_stages else 0): + if last_output: + for d in T.Pipelined(4, num_stages=2): + with T.Stage(0): + T.dual_copy(cpv[d, :, :], out, unit_flag_ctrl=3) + if emitted > 0: + final_output_vf(True, d, denorm, out, accum, final) + else: + final_output_vf(False, d, denorm, out, accum, final) + T.copy(final[:, :512], O[qidx, sid * 32 : sid * 32 + 32, :]) + if pipeline_stages: + with T.Stage(q_lead_shift + (0 if extra_metadata_lead else 1) if pipeline_stages else 0): + gather_kv( + k, + job, + sid, + extra, + KV, + ExtraKV, + group_smaller, + group_larger, + None if direct_query_masks else smaller, + None if direct_query_masks else larger, + kv, + raw8, + raw4, + ) + with T.Stage(4 + q_lead_shift): + # Source order dispatches gather DMA before the scalar + # wait for softmax's delta, despite the higher stage. + ss[sid] = T.Cast("uint32", delta[0] > 6.0 / scale) + + if decode: + + @T.prim_func + def main( + Q: T.StridedTensor((nq, 64, 512), q_strides, "bfloat16"), + KV: T.Tensor((kv_storage_bytes,), "uint8"), + Indices: T.StridedTensor((nq, topk), (index_stride, 1), "int32"), + Sink: T.Tensor((64,), "float32"), + Lengths: T.Tensor((nq,), "int32"), + ExtraKV: T.Tensor((extra_storage_bytes,), "uint8"), + ExtraIndices: T.StridedTensor((nq, max(extra_topk, 1)), (extra_index_stride, 1), "int32"), + ExtraLengths: T.Tensor((nq,), "int32"), + O: T.Tensor((nq, 64, 512), "bfloat16"), + LSE: T.Tensor((nq, 64), "float32"), + ): + attention_body(Q, KV, Indices, Sink, Lengths, ExtraKV, ExtraIndices, ExtraLengths, O, None, LSE) + + else: + + @T.prim_func + def main( + Q: T.StridedTensor((nq, 64, 512), q_strides, "bfloat16"), + KV: T.Tensor((nk, 512), "bfloat16"), + Indices: T.StridedTensor((nq, topk), (index_stride, 1), "int32"), + Sink: T.Tensor((64,), "float32"), + Lengths: T.Tensor((nq,), "int32"), + O: T.Tensor((nq, 64, 512), "bfloat16"), + Max: T.Tensor((nq, 64), "float32"), + LSE: T.Tensor((nq, 64), "float32"), + ): + attention_body(Q, KV, Indices, Sink, Lengths, None, None, None, O, Max, LSE) + + return main + + +def prefill(nq, nk, topk, **kwargs): + return sparse_attention(nq, nk, topk, **kwargs) + + +def compile_prefill(nq, nk, topk, **kwargs): + # UB ND-to-NZ copy layout inference currently needs an active target context. + with tvm.target.Target("ascend"): + return tilelang.compile( + prefill(nq, nk, topk, **kwargs), + target="ascend", + out_idx=[5, 6, 7], + pass_configs=PASS_CONFIG, + compile_flags=["-Ofast", "-mllvm", "-enable-hiipu-vf-loop-unroll"], + ) diff --git a/examples/ascend/flashmla/test_staged.py b/examples/ascend/flashmla/test_staged.py new file mode 100644 index 0000000000..d225cab45b --- /dev/null +++ b/examples/ascend/flashmla/test_staged.py @@ -0,0 +1,240 @@ +"""Device regression for automatic versions across rescale, jobs, and launches.""" + +import math + +import pytest +import torch + +from examples.ascend.flashmla.staged import compile_prefill + + +@pytest.mark.parametrize("pipeline_stages", [True, False], ids=["pipeline", "flat"]) +def test_last_tile_rescale_across_jobs_and_launches(pipeline_stages): + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Ascend NPU required") + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + nq, nk, topk = 33, 128, 128 + q = torch.zeros((nq, 64, 512), dtype=torch.bfloat16, device="npu") + q[:, 32:, :] = 1 + kv = torch.zeros((nk, 512), dtype=torch.bfloat16, device="npu") + indices = torch.arange(topk, dtype=torch.int32, device="npu").expand(nq, topk).contiguous() + sink = torch.zeros(64, dtype=torch.float32, device="npu") + lengths = torch.full((nq,), topk, dtype=torch.int32, device="npu") + kernel = compile_prefill(nq, nk, topk, sink=False, scale=0.5, pipeline_stages=pipeline_stages) + + # Only AIV1 requests rescaling, on the last KV tile. The second query job + # reuses all storage; changing KV between launches detects stale state. + for value in (1, 2, 1, 2, 1): + kv[64:] = value + out, maximum, lse = kernel(q, kv, indices, sink, lengths) + torch.npu.synchronize() + expected = torch.full((nq, 64, 512), value / 2, dtype=torch.bfloat16) + expected[:, 32:] = value + expected_max = torch.zeros((nq, 64), dtype=torch.float32) + expected_max[:, 32:] = 256 * value + expected_lse = torch.full((nq, 64), math.log(128), dtype=torch.float32) + expected_lse[:, 32:] = 256 * value + math.log(64) + torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(maximum.cpu(), expected_max, rtol=0, atol=0) + torch.testing.assert_close(lse.cpu(), expected_lse, rtol=0, atol=2e-5) + + +def test_index_group_tail_and_query_state(): + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Ascend NPU required") + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + # Eleven KV tiles cross the 640-index group boundary and rotate versions + # by an odd count. Block 0 processes three distinct queries. + nq, nk, topk = 65, 704, 704 + q = torch.zeros((nq, 64, 512), dtype=torch.bfloat16, device="npu") + row_values = (torch.arange(nk, device="npu") % 7).to(torch.bfloat16) + kv = row_values[:, None].expand(nk, 512).contiguous() + indices = (torch.arange(topk)[None, :] + torch.arange(nq)[:, None] * 11) % nk + indices[:, ::17] = -1 + indices = indices.to(device="npu", dtype=torch.int32) + lengths = torch.tensor([704, 640, 64, 0], dtype=torch.int32).repeat(17)[:nq].to("npu") + sink = torch.zeros(64, dtype=torch.float32, device="npu") + kernel = compile_prefill(nq, nk, topk, sink=False, variable_lengths=True) + out, maximum, lse = kernel(q, kv, indices, sink, lengths) + torch.npu.synchronize() + + indices_cpu, lengths_cpu = indices.cpu(), lengths.cpu() + valid = (indices_cpu >= 0) & (torch.arange(topk)[None, :] < lengths_cpu[:, None]) + count = valid.sum(dim=1) + values = row_values.cpu()[indices_cpu.clamp(min=0)].float() + average = (values * valid).sum(dim=1) / count.clamp(min=1) + expected = average.to(torch.bfloat16)[:, None, None].expand(nq, 64, 512) + expected_max = torch.where(count > 0, 0.0, -torch.inf)[:, None].expand(nq, 64) + expected_lse = torch.where(count > 0, count.float().log(), torch.inf)[:, None].expand(nq, 64) + torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0.015625) + torch.testing.assert_close(maximum.cpu(), expected_max, rtol=0, atol=0) + torch.testing.assert_close(lse.cpu(), expected_lse, rtol=0, atol=2e-5) + + +def test_alternating_rescale_and_accumulate(): + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Ascend NPU required") + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + nq, nk = 65, 448 + q = torch.zeros((nq, 64, 512), dtype=torch.bfloat16, device="npu") + q[:, 32:, :] = 1 + q[1::2, 32:, :] = -1 + tile_values = torch.tensor([0, 1, 1, 2, 2, 3, 3], dtype=torch.bfloat16) + values = tile_values.repeat_interleave(64) + kv = values[:, None].expand(nk, 512).contiguous().to("npu") + indices = torch.arange(nk, dtype=torch.int32).expand(nq, nk).clone() + indices[2::4] = indices[2::4].flip(1) + indices[3::4] = indices[3::4].flip(1) + indices = indices.to("npu") + sink = torch.zeros(64, dtype=torch.float32, device="npu") + lengths = torch.full((nq,), nk, dtype=torch.int32, device="npu") + kernel = compile_prefill(nq, nk, nk, sink=False, scale=0.5) + + # Alternating large score jumps exercise conditional handoffs; reversed + # queries swap which AIV1 rows rescale. Odd tile count rotates ring slots. + expected = torch.full((nq, 64, 512), 12 / 7, dtype=torch.bfloat16) + expected[::2, 32:, :] = 3 + expected[1::2, 32:, :] = 0 + expected_max = torch.zeros((nq, 64), dtype=torch.float32) + expected_max[::2, 32:] = 768 + expected_lse = torch.full((nq, 64), math.log(nk), dtype=torch.float32) + expected_lse[::2, 32:] = 768 + math.log(128) + expected_lse[1::2, 32:] = math.log(64) + for _ in range(3): + out, maximum, lse = kernel(q, kv, indices, sink, lengths) + torch.npu.synchronize() + torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0.0078125) + torch.testing.assert_close(maximum.cpu(), expected_max, rtol=0, atol=0) + torch.testing.assert_close(lse.cpu(), expected_lse, rtol=0, atol=8e-5) + + +@pytest.mark.parametrize("nb", [5, 15, 16], ids=["metadata-lead-start", "q2-limit", "single-q-fallback"]) +def test_query_prefetch_boundaries(nb): + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Ascend NPU required") + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + # Three query jobs force the two-Q path to reuse its first physical slot. + nq, nk, topk = 65, nb * 64, nb * 64 + q = torch.zeros((nq, 64, 512), dtype=torch.bfloat16, device="npu") + row_values = (torch.arange(nk, device="npu") % 7).to(torch.bfloat16) + kv = row_values[:, None].expand(nk, 512).contiguous() + indices = (torch.arange(topk)[None, :] + torch.arange(nq)[:, None] * 11) % nk + indices[:, ::17] = -1 + indices = indices.to(device="npu", dtype=torch.int32) + lengths = torch.tensor([0, 1, 63, 64, 65, topk - 1, topk], dtype=torch.int32).repeat(10)[:nq].to("npu") + sink = torch.zeros(64, dtype=torch.float32, device="npu") + kernel = compile_prefill(nq, nk, topk, sink=False, variable_lengths=True) + out, maximum, lse = kernel(q, kv, indices, sink, lengths) + torch.npu.synchronize() + + indices_cpu, lengths_cpu = indices.cpu(), lengths.cpu() + valid = (indices_cpu >= 0) & (torch.arange(topk)[None, :] < lengths_cpu[:, None]) + count = valid.sum(dim=1) + values = row_values.cpu()[indices_cpu.clamp(min=0)].float() + average = (values * valid).sum(dim=1) / count.clamp(min=1) + expected = average.to(torch.bfloat16)[:, None, None].expand(nq, 64, 512) + expected_max = torch.where(count > 0, 0.0, -torch.inf)[:, None].expand(nq, 64) + expected_lse = torch.where(count > 0, count.float().log(), torch.inf)[:, None].expand(nq, 64) + torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0.015625) + torch.testing.assert_close(maximum.cpu(), expected_max, rtol=0, atol=0) + torch.testing.assert_close(lse.cpu(), expected_lse, rtol=0, atol=2e-5) + + +@pytest.fixture +def flashmla_reference(monkeypatch): + """The same separately built reference used by run_decode.py.""" + import importlib + import os + from pathlib import Path + + location = os.environ.get("FLASHMLA_REFERENCE") + if not location: + pytest.skip("Set FLASHMLA_REFERENCE to the pinned FlashMLA reference checkout") + pytest.importorskip("torch_npu") + if not torch.npu.is_available(): + pytest.skip("Ascend NPU required") + reference = Path(location).resolve() + monkeypatch.syspath_prepend(str(reference / "tests")) + monkeypatch.syspath_prepend(str(reference)) + modules = tuple(importlib.import_module(name) for name in ("lib", "ref", "quant", "kernelkit")) + for module in modules: + assert Path(module.__file__).resolve().is_relative_to(reference), f"Unexpected reference module: {module.__file__}" + return modules + + +@pytest.mark.parametrize("extra_format", ["fp8", "fp4"]) +@pytest.mark.parametrize("main_lengths", [True, False], ids=["main-lengths", "extra-lengths"]) +def test_decode_independent_lengths_and_group_tails(flashmla_reference, extra_format, main_lengths): + from dataclasses import replace + + from examples.ascend.flashmla.decode import prepare_decode + + lib, ref, quant, kk = flashmla_reference + torch.npu.set_device(0) + torch.npu.set_op_timeout_ms(30000) + fp8 = quant.KVCacheLayout.V41_FP8Sparse + fmt = quant.KVCacheLayout.V41_FP4Sparse if extra_format == "fp4" else fp8 + rp = lib.RawTestParamForDecode( + 35, + 64, + 3, + 1, + 1153, + False, + 704, + extra_s_k=1537, + extra_topk=768, + block_size=53, + extra_block_size=61, + have_topk_length=main_lengths, + have_extra_topk_length=not main_lengths, + have_zero_seqlen_k=False, + enable_attn_sink=True, + kvcache_layout=fp8, + extra_kvcache_layout=fmt, + seed=42, + ) + p = rp.to_test_param() + # Generate every selected KV entry before imposing boundary lengths. + # Otherwise the reference fixture poisons entries beyond random lengths. + fixture_param = replace(p, have_topk_length=False, decode=replace(p.decode, have_extra_topk_length=False)) + old_dtype = torch.get_default_dtype() + old_precision = torch.get_float32_matmul_precision() + try: + torch.set_default_dtype(torch.bfloat16) + torch.set_float32_matmul_precision("high") + with torch.device("npu:0"): + inputs = lib.generate_testcase_for_decode(fixture_param) + inputs.p = p + inputs.kv_scope.t = p + inputs.extra_kv_scope.t = p + max_length = rp.topk if rp.have_topk_length else rp.extra_topk + boundary_lengths = [0, 1, 63, 64, 65, 639, 640, 641, max_length] + pattern = torch.tensor(boundary_lengths, dtype=torch.int32, device=inputs.q.device) + active_lengths = pattern.repeat((rp.b + len(boundary_lengths) - 1) // len(boundary_lengths))[: rp.b].contiguous() + if rp.have_topk_length: + inputs.kv_scope.topk_length = active_lengths + else: + inputs.extra_kv_scope.topk_length = active_lengths + assert (inputs.kv_scope.topk_length is not None) != (inputs.extra_kv_scope.topk_length is not None) + assert not inputs.q.is_contiguous() + assert not inputs.kv_scope.indices_in_kvcache.is_contiguous() + assert not inputs.extra_kv_scope.indices_in_kvcache.is_contiguous() + torch.npu.synchronize() + _, run = prepare_decode(p, inputs, implementation="staged") + out, lse = run() + torch.npu.synchronize() + expected = ref.ref_sparse_attn_decode(p, inputs) + assert kk.check_is_allclose("staged out", out, expected[0], abs_tol=1e-3, rel_tol=2.01 / 128, cos_diff_tol=5e-6) + assert kk.check_is_allclose("staged lse", lse, expected[1], abs_tol=1e-6, rel_tol=8.01 / 65536) + finally: + torch.set_default_dtype(old_dtype) + torch.set_float32_matmul_precision(old_precision) From 34dff840e628b6483aa75e09783757520c22b850 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:07:21 +0800 Subject: [PATCH 02/12] [Ascend] Link the separate FlashMLA compiler prerequisites --- examples/ascend/flashmla/WORKTREES.md | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/examples/ascend/flashmla/WORKTREES.md b/examples/ascend/flashmla/WORKTREES.md index 90b15da16c..a0ed94c3bd 100644 --- a/examples/ascend/flashmla/WORKTREES.md +++ b/examples/ascend/flashmla/WORKTREES.md @@ -6,16 +6,16 @@ synchronized to upstream `d85a70a97218fec25856f8b5c969f167ad96dc1e` on 2026-10-2 Apply the compiler prerequisites before building the staged kernel. The source does not use `T.IssuePhase`, and its API is excluded from these proposals. -| Compiler change | Fork branch | -| --- | --- | -| Support native NZ output for L0C-to-UB copies | [codex/ascend-l0c-ub-nz](https://github.com/Denverjin/tilelang/tree/codex/ascend-l0c-ub-nz) | -| Preserve Cube configuration phases across manual stages | [codex/ascend-cube-register-phases](https://github.com/Denverjin/tilelang/tree/codex/ascend-cube-register-phases) | -| Select AIV streams outside scheduled pipelines | [codex/ascend-aiv-stream-selection](https://github.com/Denverjin/tilelang/tree/codex/ascend-aiv-stream-selection) | -| Preserve complete periodic buffer lifetimes | [codex/ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/codex/ascend-periodic-buffer-lifetimes) | -| Establish buffer placement hints from exact L1 offsets | [codex/ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/codex/ascend-buffer-placement-hints) | -| Preserve physical views and pointer units during layout lowering | [codex/ascend-layout-pointer-units](https://github.com/Denverjin/tilelang/tree/codex/ascend-layout-pointer-units) | -| Reuse storage and L0B rings for explicit conflict hints | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | -| Preserve manual source order across physical phase ties | [codex/ascend-manual-pipe-order](https://github.com/Denverjin/tilelang/tree/codex/ascend-manual-pipe-order) | +| Compiler change | Pull request | Fork branch | +| --- | --- | --- | +| Support native NZ output for L0C-to-UB copies | [#5](https://github.com/Denverjin/tilelang/pull/5) | [codex/ascend-l0c-ub-nz](https://github.com/Denverjin/tilelang/tree/codex/ascend-l0c-ub-nz) | +| Preserve Cube configuration phases across manual stages | [#6](https://github.com/Denverjin/tilelang/pull/6) | [codex/ascend-cube-register-phases](https://github.com/Denverjin/tilelang/tree/codex/ascend-cube-register-phases) | +| Select AIV streams outside scheduled pipelines | [#7](https://github.com/Denverjin/tilelang/pull/7) | [codex/ascend-aiv-stream-selection](https://github.com/Denverjin/tilelang/tree/codex/ascend-aiv-stream-selection) | +| Preserve complete periodic buffer lifetimes | [#8](https://github.com/Denverjin/tilelang/pull/8) | [codex/ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/codex/ascend-periodic-buffer-lifetimes) | +| Establish buffer placement hints from exact L1 offsets | [#9](https://github.com/Denverjin/tilelang/pull/9) | [codex/ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/codex/ascend-buffer-placement-hints) | +| Preserve physical views and pointer units during layout lowering | [#10](https://github.com/Denverjin/tilelang/pull/10) | [codex/ascend-layout-pointer-units](https://github.com/Denverjin/tilelang/tree/codex/ascend-layout-pointer-units) | +| Reuse storage and L0B rings for explicit conflict hints | [#11](https://github.com/Denverjin/tilelang/pull/11) | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | +| Preserve manual source order across physical phase ties | [#12](https://github.com/Denverjin/tilelang/pull/12) | [codex/ascend-manual-pipe-order](https://github.com/Denverjin/tilelang/tree/codex/ascend-manual-pipe-order) | The additional NZ pitch, compact-allocation, and access-footprint prerequisite is already proposed upstream as [tile-ai/tilelang#3530](https://github.com/tile-ai/tilelang/pull/3530). From e5fe270d0cf0d53713eefb0a05c18e3c2d10d581 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 15:31:52 +0800 Subject: [PATCH 03/12] [Ascend] Share FlashMLA L0B storage through layout views --- examples/ascend/flashmla/README.md | 10 ++++++---- examples/ascend/flashmla/WORKTREES.md | 15 +++++++++------ examples/ascend/flashmla/staged.py | 5 ++--- 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/examples/ascend/flashmla/README.md b/examples/ascend/flashmla/README.md index e1976c0d7d..b6cca7cf3a 100644 --- a/examples/ascend/flashmla/README.md +++ b/examples/ascend/flashmla/README.md @@ -21,10 +21,12 @@ explicit query slots when the periodic lifetime fits the event budget (`4 <= nb <= 15`); other shapes retain the single-query-slot path. The two 32-row KV halves have explicit L1 offsets in separate bank groups. -QK and PV use separate four-version L0B buffers with MajorK and MajorMN -layouts. `T.assume_conflict(bq, bv)` lets the compiler give this compatible pair -one storage-ring clock and a shared 64 KiB allocation. Their copies and matrix -products use `T.copy` and `T.gemm`. PV dispatches its nonzero and zero unit flags +QK and PV use two views of one four-version L0B allocation. QK's `bq` has a +MajorK layout, while `bv = T.view(bq, shape=(64, 128))` has a MajorMN layout. +The views share `.data`, one version clock, and a 64 KiB ring, so no +`T.assume_conflict` hint is needed. Each producer writes the layout consumed +through its view; creating the view does not convert existing data. Their copies +and matrix products use `T.copy` and `T.gemm`. PV dispatches its nonzero and zero unit flags through two GEMM branches in one `T.Task`; the two-KV-tile path keeps the flag in a scalar variable so rescale synchronization survives constant folding. diff --git a/examples/ascend/flashmla/WORKTREES.md b/examples/ascend/flashmla/WORKTREES.md index a0ed94c3bd..e7134aaefd 100644 --- a/examples/ascend/flashmla/WORKTREES.md +++ b/examples/ascend/flashmla/WORKTREES.md @@ -14,7 +14,7 @@ The source does not use `T.IssuePhase`, and its API is excluded from these propo | Preserve complete periodic buffer lifetimes | [#8](https://github.com/Denverjin/tilelang/pull/8) | [codex/ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/codex/ascend-periodic-buffer-lifetimes) | | Establish buffer placement hints from exact L1 offsets | [#9](https://github.com/Denverjin/tilelang/pull/9) | [codex/ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/codex/ascend-buffer-placement-hints) | | Preserve physical views and pointer units during layout lowering | [#10](https://github.com/Denverjin/tilelang/pull/10) | [codex/ascend-layout-pointer-units](https://github.com/Denverjin/tilelang/tree/codex/ascend-layout-pointer-units) | -| Reuse storage and L0B rings for explicit conflict hints | [#11](https://github.com/Denverjin/tilelang/pull/11) | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | +| Preserve independent Cube view layouts on shared storage rings | [#11](https://github.com/Denverjin/tilelang/pull/11) | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | | Preserve manual source order across physical phase ties | [#12](https://github.com/Denverjin/tilelang/pull/12) | [codex/ascend-manual-pipe-order](https://github.com/Denverjin/tilelang/tree/codex/ascend-manual-pipe-order) | The additional NZ pitch, compact-allocation, and access-footprint prerequisite @@ -27,15 +27,18 @@ Its PR describes the intended extension to manual placement hints for every buffer kind. That extension is not implemented yet; the kernel still calls `T.annotate_l1_offsets` and uses the existing L1 placement contract. -The conflict-hint changes have two parts: permission for pairwise storage reuse -and a shared version clock/ownership protocol for a compatible isolated L0B pair. -Both are needed for the separate QK/PV layouts to share a physical 64 KiB ring. +The staged kernel requires PR #11's Cube-view layout support from commit +`393a0bde46d60a3b7cfaabeeb65ebc67a7a14ec7` or later. QK allocates one +four-version L0B buffer; PV uses a same-shape `T.view` with a different layout. +The shared `.data` gives both views one version clock and a physical 64 KiB +ring without an `assume_conflict` hint. Both views use BF16 and the same padded +physical size. Copy and GEMM lowering preserve the selected view's layout. The raw FP8 and FP4 buffers retain independent storage and three-slot histories. ## Historical integrated measurements -The implementation source in this branch is unchanged from the integrated -kernel at `00e7b9bb64c892c5323f42bede764e8940ceb4cd`. +The historical measurements below used the integrated kernel at +`00e7b9bb64c892c5323f42bede764e8940ceb4cd`, before the L0B view change. That compiler was based on upstream main `876a59e`; its documentation was committed as `6bdfe6e`. The historical integration still contained an unused IssuePhase API. These identities describe the recorded experiments, not the diff --git a/examples/ascend/flashmla/staged.py b/examples/ascend/flashmla/staged.py index ee9e49b7c4..d98e1c3332 100644 --- a/examples/ascend/flashmla/staged.py +++ b/examples/ascend/flashmla/staged.py @@ -353,7 +353,8 @@ def attention_body(Q, KV, Indices, Sink, Lengths, ExtraKV, ExtraIndices, ExtraLe s_l1 = T.alloc_l1((64, 64), "bfloat16") aq = T.alloc_l0a((64, 128), "bfloat16") bq = T.alloc_l0b((64, 128), "bfloat16") - bv = T.alloc_l0b((64, 128), "bfloat16") + # QK and PV keep distinct layouts on one physical versioned ring. + bv = T.view(bq, shape=(64, 128)) cq = T.alloc_l0c((64, 64), "float32") # Four disjoint output-column tiles, each accumulated across all KV tiles. cpv = T.alloc_l0c((4, 64, 128), "float32") @@ -418,7 +419,6 @@ def attention_body(Q, KV, Indices, Sink, Lengths, ExtraKV, ExtraIndices, ExtraLe s_l1: 2, aq: 4, bq: 4, - bv: 4, cq: 1, cpv: 1, kv: 1 if decode else 3, @@ -469,7 +469,6 @@ def attention_body(Q, KV, Indices, Sink, Lengths, ExtraKV, ExtraIndices, ExtraLe # Each nested scheduler loop is manual independently. VF loops # belong to their enclosing VF task and cannot contain T.Stage. for step in T.Pipelined(T.ceildiv(nq - bx, 32) * nb, num_stages=6 + q_lead_shift, annotations={"enable_offset": True}): - T.assume_conflict(bq, bv) job = step // nb k = step % nb qidx = bx + job * 32 From c12b4383009b942623f2e47fb297436ed44760d6 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 14:57:20 +0800 Subject: [PATCH 04/12] [Ascend] Support NZ layouts for L0C-to-UB copies (cherry picked from commit 33f29b01b7473c0a3d34643148683a69e481074a) --- src/ascend/op/copy.cc | 50 ++++++++++++- .../transform/normalize_fractal_storage.cc | 5 ++ ...test_ascend_fractal_storage_integration.py | 75 +++++++++++++++++++ ...est_tilelang_ascend_fractal_layout_tags.py | 5 ++ .../transform/test_ascend_lower_copy.py | 45 +++++++++++ tilelang/layout/__init__.py | 1 + tilelang/layout/ascend.py | 21 ++++++ 7 files changed, 199 insertions(+), 3 deletions(-) diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 302d832b99..8025d66abc 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -663,8 +663,15 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, } else if (dma_path == DMAPath::kL0CToUB) { StridedLayout src_layout = NormalizeTrailingMTE2DLayout( op.src, src_range, analyzer, "Ascend L0C->UB source"); - StridedLayout dst_layout = NormalizeMTE2DLayout( - op.dst, dst_range, analyzer, "Ascend L0C->UB destination"); + Optional ub_layout = FindLayoutForBuffer(T.layout_map, op.dst); + bool nz_output = + ub_layout.defined() && IsAscendFractalLayout(ub_layout.value()); + StridedLayout dst_layout = + nz_output + ? NormalizeTrailingMTE2DLayout(op.dst, dst_range, analyzer, + "Ascend NZ L0C->UB destination") + : NormalizeMTE2DLayout(op.dst, dst_range, analyzer, + "Ascend L0C->UB destination"); const Mode &src_inner = src_layout.modes[0]; const Mode &src_row = src_layout.modes[1]; const Mode &dst_inner = dst_layout.modes[0]; @@ -721,6 +728,43 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, copy_inner = dst_inner.size; copy_rows = dst_row.size; } + if (nz_output) { + AscendFractalLayoutInfo dst_info; + ICHECK( + TryExtractAscendFractalLayout(ub_layout.value(), op.dst, &dst_info)); + TVM_FFI_CHECK(dst_info.c0_axis == LogicalAxis::kCol && + dst_info.row16_axis == LogicalAxis::kRow && + analyzer->CanProveEqual(dst_info.c0, 16), + ValueError) + << "Ascend NZ L0C->UB requires 16-column NZ panels on " + << op.dst->name << "; use make_ascend_ub_nz_layout."; + if (!op.dst->strides.empty()) { + PrimExpr stride = I(1); + for (int i = static_cast(op.dst->shape.size()) - 1; i >= 0; --i) { + TVM_FFI_CHECK(analyzer->CanProveEqual(op.dst->shape[i], 1) || + analyzer->CanProveEqual(op.dst->strides[i], stride), + ValueError) + << "Ascend NZ L0C->UB requires dense logical strides; express " + "padding in the UB buffer shape instead of explicit strides."; + stride = stride * op.dst->shape[i]; + } + } + With nonempty_copy(analyzer, has_data); + TVM_FFI_CHECK( + analyzer->CanProveEqual(FloorMod(dst_range.back()->min, I(16)), 0) && + analyzer->CanProveEqual(FloorMod(dst_inner.size, I(16)), 0) && + analyzer->CanProveEqual(FloorMod(copy_inner, I(16)), 0), + ValueError) + << "Ascend NZ L0C->UB requires a destination column origin and " + "copy widths aligned to 16 elements; got " + << dst_range.back() << " and source width " << copy_inner << "."; + // NZ pitch advances between 16-column panels, in destination elements. + // Use the allocation's physical shape, not this copy's row extent. + Array physical_shape = ub_layout.value()->OutputShape(); + dst_row_stride = I(1); + for (size_t i = physical_shape.size() - 3; i < physical_shape.size(); ++i) + dst_row_stride = dst_row_stride * physical_shape[i]; + } PrimExpr dual_dst_ctl_val = I(dual_dst_ctl); PrimExpr unit_flag_ctl_val = cast(DataType::Int(32), op.unit_flag_ctl); PrimExpr sub_blockid_val = cast(DataType::Int(32), op.sub_blockid); @@ -739,7 +783,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, I(GetCCQuantPre(op.src->dtype, op.dst->dtype)), // quant_pre zero, // relu_pre zero, // split_en - I(1), // NZ2ND_en + I(nz_output ? 0 : 1), // NZ2ND_en zero, // quant_post zero, // relu_post zero, // clip_relu_post diff --git a/src/ascend/transform/normalize_fractal_storage.cc b/src/ascend/transform/normalize_fractal_storage.cc index 7fb7228b82..1a19a98a01 100644 --- a/src/ascend/transform/normalize_fractal_storage.cc +++ b/src/ascend/transform/normalize_fractal_storage.cc @@ -276,6 +276,11 @@ class FractalStorageNormalizer : public arith::IRMutatorWithAnalyzer { Call NormalizeFixpipeCopy_(Call call) { constexpr int kFixpipeAlignmentBytes = 32; AscendCopy rewritten(call->args, call->annotations); + // NZ output has a panel pitch and whole-column-panel requirements, not + // ND's row pitch and 32-byte row padding. Copy lowering validates it. + if (auto layout = canonical_layouts_.Get(rewritten->dst); + layout && IsAscendFractalLayout(layout.value())) + return call; ICHECK_GE(rewritten->src_range.size(), 2U); ICHECK_EQ(rewritten->src_range.size(), rewritten->src->shape.size()); // Lowering skips the DMA if any source or destination extent is diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index f8ccf45c7c..7c93948c25 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/testing/ascend/layout/test_ascend_fractal_storage_integration.py @@ -5,10 +5,85 @@ import tilelang.ascend.language as T import tilelang.testing from tilelang import tvm +from tilelang.layout import make_ascend_ub_nz_layout from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, statements +def _nz_fixpipe_kernel(split, dtype="float32", n=64, alias=False): + rows = 16 if split == "m" else 32 + cols = n // 2 if split == "n" else n + + @T.prim_func + def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: T.Tensor((32, n), dtype)): + with T.MixedKernel(1, sids=1 if split == "single" else 2) as (_, sid): + a = T.alloc_l1((32, 16), "bfloat16") + b = T.alloc_l1((n, 16), "bfloat16") + acc = T.alloc_l0c((32, n), "float32") + if alias: + storage = T.alloc_shared((2 * 48 * 80,), dtype) + ub = T.view(storage, (2, 48, 80)) + else: + ub = T.alloc_shared((2, 48, 80), dtype) + T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)}) + T.copy(A, a) + T.copy(B, b) + T.gemm(a, b, acc, transpose_B=True, clear_accum=True) + if split == "single": + T.copy(acc, ub[1, 16 : 16 + rows, 16 : 16 + cols]) + else: + T.dual_copy(acc, ub[1, 16 : 16 + rows, 16 : 16 + cols]) + with T.SimtVF(threads=128): + for i, j in T.Parallel(rows, cols): + C[i + (sid * rows if split == "m" else 0), j + (sid * cols if split == "n" else 0)] = ub[1, 16 + i, 16 + j] + + return main + + +@pytest.mark.parametrize("alias", [False, True]) +def test_nz_fixpipe_layout_survives_full_lowering(alias): + snapshots = {} + + @tvm.ir.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name == "tl.AscendLowerTileOp": + snapshots[info.name] = mod + + with tvm.target.Target("ascend"), tvm.transform.PassContext(instruments=[Capture()]): + tilelang.lower(_nz_fixpipe_kernel("m", alias=alias), target="ascend") + (copy,) = calls(snapshots["tl.AscendLowerTileOp"], "tl.ascend_copy_matrix_cc_to_ub") + assert int(copy.args[5]) == 48 * 16 + assert int(copy.args[7]) == 1 + assert int(copy.args[14]) == 0 + + +def test_nz_fixpipe_does_not_apply_nd_tail_padding(): + with tvm.target.Target("ascend"), pytest.raises(ValueError, match="aligned to 16 elements"): + tilelang.lower(_nz_fixpipe_kernel("single", n=9), target="ascend") + + +@tilelang.testing.requires_ascend +@pytest.mark.parametrize( + "split,dtype,alias", + [("single", "float32", False), ("m", "float32", False), ("n", "float32", False), ("single", "bfloat16", False), ("m", "float32", True)], +) +def test_nz_fixpipe_logical_reads_npu(split, dtype, alias): + import torch + + # Exact small inputs make addressing errors distinguishable from rounding. + generator = torch.Generator().manual_seed(42) + a = torch.randint(-4, 5, (32, 16), generator=generator).to(torch.bfloat16) + b = torch.randint(-4, 5, (64, 16), generator=generator).to(torch.bfloat16) + with tvm.target.Target("ascend"): + kernel = tilelang.compile(_nz_fixpipe_kernel(split, dtype, alias=alias), out_idx=-1, target="ascend") + expected = (a.float() @ b.float().T).to(getattr(torch, dtype)) + for _ in range(2): + actual = kernel(a.npu(), b.npu()) + torch.npu.synchronize() + torch.testing.assert_close(actual.cpu(), expected, atol=0, rtol=0) + + def test_padded_alias_storage_sets_the_physical_version_pitch(): @T.prim_func def main(A: T.Tensor((60, 15), "bfloat16")): diff --git a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index 2e5584e20a..d12e92f523 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,6 +8,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, + make_ascend_ub_nz_layout, try_extract_fractal_layout, ) @@ -18,6 +19,7 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), + (make_ascend_ub_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -29,6 +31,9 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected + if factory == make_ascend_ub_nz_layout: + assert int(info.c0) == 16 + assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] def test_expand_preserves_sf_tag(): diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 525f4848df..3aab93d5b9 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -11,6 +11,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, + make_ascend_ub_nz_layout, ) from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, kernel @@ -159,6 +160,50 @@ def test_fix_uses_destination_extent_and_preserves_source_pitch(case): _equal(copy.args[3:7], [12, 8, 40, 16] if case == "grouped" else [16, 72, 16, 128]) +@pytest.mark.parametrize("dual,src_shape", [(0, (32, 64)), (1, (64, 64)), (2, (32, 128))], ids=["single", "split-m", "split-n"]) +def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape): + src = tirx.decl_buffer(src_shape, "float32", name="src", scope="shared.l0c") + dst = tirx.decl_buffer((2, 64, 96), "float32", name="dst", scope="shared") + after = _lower( + src, + dst, + dst_region=_region(dst, [1, 32, 64], [1, 16, 16]), + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + annotations={"dual_dst_ctl": dual}, + unit_flag_ctrl=3, + ) + (copy,) = calls(after, "tl.ascend_copy_matrix_cc_to_ub") + _equal(copy.args[3:8], [src_shape[1], src_shape[0], 64 * 16, src_shape[0], dual]) + _equal([copy.args[0].args[2], copy.args[10], copy.args[14]], [64 * 96 + 64 * 16 + 16 * 16, 3, 0]) + + +@pytest.mark.parametrize("column,width", [(8, 32), (0, 24)]) +def test_fix_nz_rejects_partial_column_panels(column, width): + src = tirx.decl_buffer((32, width), "float32", name="src", scope="shared.l0c") + dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared") + with pytest.raises(ValueError, match="aligned to 16 elements"): + _lower( + src, + dst, + dst_region=_region(dst, [32, width], [0, column]), + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + ) + + +def test_fix_nz_rejects_l1_float32_panel_width(): + src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") + dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared") + with pytest.raises(ValueError, match="use make_ascend_ub_nz_layout"): + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) + + +def test_fix_nz_rejects_noncompact_logical_strides(): + src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") + dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared", strides=(80, 1)) + with pytest.raises(ValueError, match="requires dense logical strides"): + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}) + + @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"]) def test_nonrepresentable_dma_geometry_is_rejected(scope): src = tirx.decl_buffer((2, 3, 4), "bfloat16", name="src", strides=(20, 5, 1) if scope == "shared.l1" else (14, 4, 1)) diff --git a/tilelang/layout/__init__.py b/tilelang/layout/__init__.py index d13766127b..4389b72ba7 100644 --- a/tilelang/layout/__init__.py +++ b/tilelang/layout/__init__.py @@ -26,6 +26,7 @@ make_ascend_major_k_layout, # noqa: F401 make_ascend_major_mn_layout, # noqa: F401 make_ascend_l0c_layout, # noqa: F401 + make_ascend_ub_nz_layout, # noqa: F401 make_ascend_sf_layout, # noqa: F401 make_strided_slice, # noqa: F401 try_extract_fractal_layout, # noqa: F401 diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index cfcaf8560f..078d60afcc 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -92,6 +92,27 @@ def make_ascend_l0c_layout(buffer: BufferLikeType): return _ffi_api.make_ascend_l0c_layout(buf) +def make_ascend_ub_nz_layout(buffer: BufferLikeType): + """Declare the UB NZ format produced by an L0C-to-UB copy. + + Use with ``T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)})``. + ``T.copy`` and ``T.dual_copy`` then preserve NZ instead of converting to + ND. Logical ``[..., m, n]`` maps to + ``[..., n // 16, m // 16, m % 16, n % 16]``. Unlike the L1 NZ layout, + the column block has 16 elements even for float32. + + Declare any padding in the buffer shape; copy regions retain the + allocation's physical panel pitch. NZ copies require whole 16-column + panels. Ordinary indexing uses logical coordinates; raw SIMD pointers + still access consecutive physical elements. + """ + buf, _, _ = _get_buffer_info(buffer) + if buf.scope() not in ("shared", "shared.dyn"): + raise ValueError("make_ascend_ub_nz_layout requires a UB buffer") + # Fixpipe preserves the accumulator's 16x16 ordering in UB. + return _ffi_api.make_ascend_l0c_layout(buf) + + def make_ascend_sf_layout(buffer: BufferLikeType): """ Create an Ascend SF layout for MX scale-factor L1 buffers. From 9ea0d961b645e9c205407e55e109db8fba35f930 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 15:11:58 +0800 Subject: [PATCH 05/12] [Ascend] Select NZ layouts from buffer scope (cherry picked from commit e1a13a6f21303765954ad77acc8a05f06bc253b6) --- src/ascend/layout/ascend_layouts.cc | 8 +++- src/ascend/layout/ascend_layouts.h | 5 ++- src/ascend/op/copy.cc | 2 +- src/ascend/transform/insert_nd2nz.cc | 5 +++ ...test_ascend_fractal_storage_integration.py | 4 +- ...est_tilelang_ascend_fractal_layout_tags.py | 17 +++++-- .../transform/test_ascend_insert_nd2nz.py | 25 ++++++++++- .../transform/test_ascend_lower_copy.py | 11 +++-- tilelang/layout/__init__.py | 1 - tilelang/layout/ascend.py | 44 +++++++------------ 10 files changed, 76 insertions(+), 46 deletions(-) diff --git a/src/ascend/layout/ascend_layouts.cc b/src/ascend/layout/ascend_layouts.cc index f141c20e70..489a2f9bbc 100644 --- a/src/ascend/layout/ascend_layouts.cc +++ b/src/ascend/layout/ascend_layouts.cc @@ -38,8 +38,8 @@ static Layout ExpandLayout2D(const Layout &base, const Buffer &buffer) { // ND2NZ convention where the contiguous axis (cols) is grouped by C0 and is the // outer physical major. Within a fractal the row lane remains stride C0 and the // column lane remains stride 1. MajorK means cols is the reduce-K axis; MajorMN -// means cols is the output M/N axis. The legacy NZ API remains an alias of -// MajorK while the frontend migrates to explicit major names. +// means cols is the output M/N axis. The NZ API selects this layout for Cube +// inputs and the fixed accumulator layout for L0C and UB Fixpipe output. static constexpr int kAscendRowFractal = 16; namespace { @@ -198,6 +198,10 @@ Layout MakeAscendMajorMNLayout(const Buffer &buffer, int k_align) { } Layout MakeAscendNZLayout(const Buffer &buffer) { + ICHECK(buffer.defined()) << "NZ layout expects a defined buffer"; + const String scope = buffer.scope(); + if (scope == "shared" || scope == "shared.dyn" || scope == "shared.l0c") + return MakeAscendL0CLayout(buffer); return MakeAscendMajorKLayout(buffer); } diff --git a/src/ascend/layout/ascend_layouts.h b/src/ascend/layout/ascend_layouts.h index 24d731191a..4f93c125ad 100644 --- a/src/ascend/layout/ascend_layouts.h +++ b/src/ascend/layout/ascend_layouts.h @@ -18,12 +18,13 @@ namespace tl { // tile maps to physical [.., cols / C0, rows / 16, 16, C0] (a grid of 16xC0 // fractals, row-major within a fractal) where C0 = 32 bytes / element_size // (16 for bf16/fp16, 32 for fp8, 8 for fp32). In MajorK, cols is the reduce-K -// axis; in MajorMN, cols is the output M/N axis. The legacy NZ name is kept as -// an alias of MajorK during the transition. +// axis; in MajorMN, cols is the output M/N axis. Layout MakeAscendMajorKLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendMajorMNLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendL0CLayout(const tirx::Buffer &buffer); Layout MakeAscendSFLayout(const tirx::Buffer &buffer); +// NZ selects the fixed 16x16 accumulator layout for L0C and UB Fixpipe +// output, and MajorK for other scopes. Compact UB packing remains explicit. Layout MakeAscendNZLayout(const tirx::Buffer &buffer); // Compact padded-NZ layout used by software UB ND->NZ scatter: diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index 8025d66abc..a8bd518002 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -737,7 +737,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, analyzer->CanProveEqual(dst_info.c0, 16), ValueError) << "Ascend NZ L0C->UB requires 16-column NZ panels on " - << op.dst->name << "; use make_ascend_ub_nz_layout."; + << op.dst->name << "; use make_ascend_nz_layout."; if (!op.dst->strides.empty()) { PrimExpr stride = I(1); for (int i = static_cast(op.dst->shape.size()) - 1; i >= 0; --i) { diff --git a/src/ascend/transform/insert_nd2nz.cc b/src/ascend/transform/insert_nd2nz.cc index a4e7ec9c5d..c38d83cddc 100644 --- a/src/ascend/transform/insert_nd2nz.cc +++ b/src/ascend/transform/insert_nd2nz.cc @@ -227,6 +227,11 @@ class Nd2NzInserter : public arith::IRMutatorWithAnalyzer { Array ValidatePrepackedNZSource(const CopyNode ©, PrimExpr rows, PrimExpr cols, PrimExpr full_rows) const { + TVM_FFI_CHECK( + IsAscendCompactNZLayout(LayoutFor(copy.src).value(), copy.src), + ValueError) + << "UB-to-L1 NZ post-copy requires compact NZ storage; use " + "make_ascend_compact_nz_layout for UB input tiles."; RequireStaticInt(rows, "NZ post-copy rows"); RequireStaticInt(cols, "NZ post-copy cols"); RequireStaticInt(full_rows, "NZ post-copy full_rows"); diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index 7c93948c25..06862d983a 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/testing/ascend/layout/test_ascend_fractal_storage_integration.py @@ -5,7 +5,7 @@ import tilelang.ascend.language as T import tilelang.testing from tilelang import tvm -from tilelang.layout import make_ascend_ub_nz_layout +from tilelang.layout import make_ascend_nz_layout from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, statements @@ -25,7 +25,7 @@ def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: ub = T.view(storage, (2, 48, 80)) else: ub = T.alloc_shared((2, 48, 80), dtype) - T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)}) + T.annotate_layout({ub: make_ascend_nz_layout(ub)}) T.copy(A, a) T.copy(B, b) T.gemm(a, b, acc, transpose_B=True, clear_accum=True) diff --git a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index d12e92f523..8692dee50a 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,7 +8,7 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_ub_nz_layout, + make_ascend_nz_layout, try_extract_fractal_layout, ) @@ -19,7 +19,7 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), - (make_ascend_ub_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), + (make_ascend_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -31,11 +31,22 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected - if factory == make_ascend_ub_nz_layout: + if factory == make_ascend_nz_layout: assert int(info.c0) == 16 assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] +@pytest.mark.parametrize("scope,kind", [("shared", 1), ("shared.dyn", 1), ("shared.l0c", 1), ("shared.l1", 0)]) +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +def test_nz_layout_uses_buffer_scope(scope, kind, dtype): + buf = tirx.decl_buffer((2, 32, 64), dtype, name="A", scope=scope) + layout = make_ascend_nz_layout(buf) + info = try_extract_fractal_layout(layout, buf) + c0 = 8 if scope == "shared.l1" and dtype == "float32" else 16 + assert (info.kind, int(info.c0), info.c0_axis, info.row16_axis) == (kind, c0, 1, 0) + assert list(layout.get_output_shape()) == [2, 64 // c0, 2, 16, c0] + + def test_expand_preserves_sf_tag(): base_buf = tirx.decl_buffer((16, 1), "uint16", name="SF2", scope="shared") expanded_buf = tirx.decl_buffer((2, 16, 1), "uint16", name="SF3", scope="shared") diff --git a/testing/ascend/transform/test_ascend_insert_nd2nz.py b/testing/ascend/transform/test_ascend_insert_nd2nz.py index 101173b966..3e28614bd6 100644 --- a/testing/ascend/transform/test_ascend_insert_nd2nz.py +++ b/testing/ascend/transform/test_ascend_insert_nd2nz.py @@ -3,7 +3,7 @@ import pytest from tilelang import tvm from tilelang.ascend import language as T, transform -from tilelang.layout import make_ascend_compact_nz_layout +from tilelang.layout import make_ascend_compact_nz_layout, make_ascend_nz_layout from tvm import tirx from testing.ascend._ir import calls @@ -66,3 +66,26 @@ def test_scatter_storage_requirements(kwargs, error): transform.InsertNd2Nz()(before) else: assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_scatter")) == 1 + + +@pytest.mark.parametrize("factory", [make_ascend_nz_layout, make_ascend_compact_nz_layout]) +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +def test_post_copy_requires_compact_nz_panel_pitch(factory, dtype): + src = tirx.decl_buffer((33, 128), dtype, name="src", scope="shared.dyn") + dst = tirx.decl_buffer((32, 128), dtype, name="dst", scope="shared.l1") + region = tirx.BufferRegion(src, [tvm.ir.Range(0, 32), tvm.ir.Range(0, 128)]) + root = tirx.SBlock( + [], + [], + [], + "root", + tirx.Evaluate(T.copy(region, dst)), + alloc_buffers=[src, dst], + annotations={"layout_map": {src: factory(src), dst: make_ascend_nz_layout(dst)}}, + ) + before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, root))}) + if factory == make_ascend_nz_layout: + with pytest.raises(ValueError, match="requires compact NZ storage"): + transform.InsertNd2Nz()(before) + else: + assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_post_copy")) == 1 diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index 3aab93d5b9..faa0d7522c 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -11,7 +11,6 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_ub_nz_layout, ) from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, kernel @@ -168,7 +167,7 @@ def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape): src, dst, dst_region=_region(dst, [1, 32, 64], [1, 16, 16]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, annotations={"dual_dst_ctl": dual}, unit_flag_ctrl=3, ) @@ -186,22 +185,22 @@ def test_fix_nz_rejects_partial_column_panels(column, width): src, dst, dst_region=_region(dst, [32, width], [0, column]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, ) def test_fix_nz_rejects_l1_float32_panel_width(): src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared") - with pytest.raises(ValueError, match="use make_ascend_ub_nz_layout"): - _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) + with pytest.raises(ValueError, match="use make_ascend_nz_layout"): + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_major_k_layout(dst)}) def test_fix_nz_rejects_noncompact_logical_strides(): src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared", strides=(80, 1)) with pytest.raises(ValueError, match="requires dense logical strides"): - _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_ub_nz_layout(dst)}) + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"]) diff --git a/tilelang/layout/__init__.py b/tilelang/layout/__init__.py index 4389b72ba7..d13766127b 100644 --- a/tilelang/layout/__init__.py +++ b/tilelang/layout/__init__.py @@ -26,7 +26,6 @@ make_ascend_major_k_layout, # noqa: F401 make_ascend_major_mn_layout, # noqa: F401 make_ascend_l0c_layout, # noqa: F401 - make_ascend_ub_nz_layout, # noqa: F401 make_ascend_sf_layout, # noqa: F401 make_strided_slice, # noqa: F401 try_extract_fractal_layout, # noqa: F401 diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index 078d60afcc..0316543e9e 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -18,16 +18,25 @@ class _AscendFractalLayout(Layout): def make_ascend_nz_layout(buffer: BufferLikeType): - """ - Create an Ascend NZ (zN fractal) layout for a Cube operand buffer. + """Create the Ascend NZ layout appropriate for the buffer's scope. + + Logical ``[..., m, n]`` maps to + ``[..., n // C0, m // 16, m % 16, n % C0]``. L0C and UB + (``shared``/``shared.dyn``) use the fixed 16-column accumulator format. + Other scopes use the Cube input format, with + ``C0 = 32 bytes / element_size``. + + With ``T.annotate_layout({ub: make_ascend_nz_layout(ub)})``, L0C-to-UB + ``T.copy`` and ``T.dual_copy`` preserve NZ. Declare padding in the UB + shape; copy regions retain the allocation's physical panel pitch and + must cover whole 16-column panels. Buffer indices remain logical; + SIMD pointer increments access consecutive physical elements. - A logical ``[.., rows, cols]`` tile maps to physical - ``[.., rows / 16, cols / C0, 16, C0]`` -- a grid of 16xC0 fractals -- where - ``C0 = 32 bytes / element_size`` is the Cube unit's fractal width and the - forward map is ``(.., r, c) -> [.., r / 16, c / C0, r % 16, c % C0]``. + For compact UB packing before UB-to-L1 copies, use + :func:`make_ascend_compact_nz_layout` explicitly. Args: - buffer: BufferLikeType (rank >= 2, constant last-2 dims) + buffer: BufferLikeType (rank >= 2). Returns: Layout: The NZ fractal layout for the buffer. @@ -92,27 +101,6 @@ def make_ascend_l0c_layout(buffer: BufferLikeType): return _ffi_api.make_ascend_l0c_layout(buf) -def make_ascend_ub_nz_layout(buffer: BufferLikeType): - """Declare the UB NZ format produced by an L0C-to-UB copy. - - Use with ``T.annotate_layout({ub: make_ascend_ub_nz_layout(ub)})``. - ``T.copy`` and ``T.dual_copy`` then preserve NZ instead of converting to - ND. Logical ``[..., m, n]`` maps to - ``[..., n // 16, m // 16, m % 16, n % 16]``. Unlike the L1 NZ layout, - the column block has 16 elements even for float32. - - Declare any padding in the buffer shape; copy regions retain the - allocation's physical panel pitch. NZ copies require whole 16-column - panels. Ordinary indexing uses logical coordinates; raw SIMD pointers - still access consecutive physical elements. - """ - buf, _, _ = _get_buffer_info(buffer) - if buf.scope() not in ("shared", "shared.dyn"): - raise ValueError("make_ascend_ub_nz_layout requires a UB buffer") - # Fixpipe preserves the accumulator's 16x16 ordering in UB. - return _ffi_api.make_ascend_l0c_layout(buf) - - def make_ascend_sf_layout(buffer: BufferLikeType): """ Create an Ascend SF layout for MX scale-factor L1 buffers. From 817676ec41bee16cbdb3c474e1fbc7baf9b3fa19 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Mon, 12 Oct 2026 15:28:10 +0800 Subject: [PATCH 06/12] [Ascend] Reuse explicit L0C layouts for UB NZ output (cherry picked from commit e9caa1855f39d6907332fe58d2f4096e0b088337) --- src/ascend/layout/ascend_layouts.cc | 8 ++---- src/ascend/layout/ascend_layouts.h | 5 ++-- src/ascend/op/copy.cc | 2 +- src/ascend/transform/insert_nd2nz.cc | 5 ---- ...test_ascend_fractal_storage_integration.py | 4 +-- ...est_tilelang_ascend_fractal_layout_tags.py | 16 ----------- .../transform/test_ascend_insert_nd2nz.py | 25 +---------------- .../transform/test_ascend_lower_copy.py | 8 +++--- tilelang/layout/ascend.py | 28 ++++++++----------- 9 files changed, 24 insertions(+), 77 deletions(-) diff --git a/src/ascend/layout/ascend_layouts.cc b/src/ascend/layout/ascend_layouts.cc index 489a2f9bbc..f141c20e70 100644 --- a/src/ascend/layout/ascend_layouts.cc +++ b/src/ascend/layout/ascend_layouts.cc @@ -38,8 +38,8 @@ static Layout ExpandLayout2D(const Layout &base, const Buffer &buffer) { // ND2NZ convention where the contiguous axis (cols) is grouped by C0 and is the // outer physical major. Within a fractal the row lane remains stride C0 and the // column lane remains stride 1. MajorK means cols is the reduce-K axis; MajorMN -// means cols is the output M/N axis. The NZ API selects this layout for Cube -// inputs and the fixed accumulator layout for L0C and UB Fixpipe output. +// means cols is the output M/N axis. The legacy NZ API remains an alias of +// MajorK while the frontend migrates to explicit major names. static constexpr int kAscendRowFractal = 16; namespace { @@ -198,10 +198,6 @@ Layout MakeAscendMajorMNLayout(const Buffer &buffer, int k_align) { } Layout MakeAscendNZLayout(const Buffer &buffer) { - ICHECK(buffer.defined()) << "NZ layout expects a defined buffer"; - const String scope = buffer.scope(); - if (scope == "shared" || scope == "shared.dyn" || scope == "shared.l0c") - return MakeAscendL0CLayout(buffer); return MakeAscendMajorKLayout(buffer); } diff --git a/src/ascend/layout/ascend_layouts.h b/src/ascend/layout/ascend_layouts.h index 4f93c125ad..24d731191a 100644 --- a/src/ascend/layout/ascend_layouts.h +++ b/src/ascend/layout/ascend_layouts.h @@ -18,13 +18,12 @@ namespace tl { // tile maps to physical [.., cols / C0, rows / 16, 16, C0] (a grid of 16xC0 // fractals, row-major within a fractal) where C0 = 32 bytes / element_size // (16 for bf16/fp16, 32 for fp8, 8 for fp32). In MajorK, cols is the reduce-K -// axis; in MajorMN, cols is the output M/N axis. +// axis; in MajorMN, cols is the output M/N axis. The legacy NZ name is kept as +// an alias of MajorK during the transition. Layout MakeAscendMajorKLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendMajorMNLayout(const tirx::Buffer &buffer, int k_align = 1); Layout MakeAscendL0CLayout(const tirx::Buffer &buffer); Layout MakeAscendSFLayout(const tirx::Buffer &buffer); -// NZ selects the fixed 16x16 accumulator layout for L0C and UB Fixpipe -// output, and MajorK for other scopes. Compact UB packing remains explicit. Layout MakeAscendNZLayout(const tirx::Buffer &buffer); // Compact padded-NZ layout used by software UB ND->NZ scatter: diff --git a/src/ascend/op/copy.cc b/src/ascend/op/copy.cc index a8bd518002..d3689ad7c6 100644 --- a/src/ascend/op/copy.cc +++ b/src/ascend/op/copy.cc @@ -737,7 +737,7 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T, analyzer->CanProveEqual(dst_info.c0, 16), ValueError) << "Ascend NZ L0C->UB requires 16-column NZ panels on " - << op.dst->name << "; use make_ascend_nz_layout."; + << op.dst->name << "; use make_ascend_l0c_layout."; if (!op.dst->strides.empty()) { PrimExpr stride = I(1); for (int i = static_cast(op.dst->shape.size()) - 1; i >= 0; --i) { diff --git a/src/ascend/transform/insert_nd2nz.cc b/src/ascend/transform/insert_nd2nz.cc index c38d83cddc..a4e7ec9c5d 100644 --- a/src/ascend/transform/insert_nd2nz.cc +++ b/src/ascend/transform/insert_nd2nz.cc @@ -227,11 +227,6 @@ class Nd2NzInserter : public arith::IRMutatorWithAnalyzer { Array ValidatePrepackedNZSource(const CopyNode ©, PrimExpr rows, PrimExpr cols, PrimExpr full_rows) const { - TVM_FFI_CHECK( - IsAscendCompactNZLayout(LayoutFor(copy.src).value(), copy.src), - ValueError) - << "UB-to-L1 NZ post-copy requires compact NZ storage; use " - "make_ascend_compact_nz_layout for UB input tiles."; RequireStaticInt(rows, "NZ post-copy rows"); RequireStaticInt(cols, "NZ post-copy cols"); RequireStaticInt(full_rows, "NZ post-copy full_rows"); diff --git a/testing/ascend/layout/test_ascend_fractal_storage_integration.py b/testing/ascend/layout/test_ascend_fractal_storage_integration.py index 06862d983a..c7c04820ba 100644 --- a/testing/ascend/layout/test_ascend_fractal_storage_integration.py +++ b/testing/ascend/layout/test_ascend_fractal_storage_integration.py @@ -5,7 +5,7 @@ import tilelang.ascend.language as T import tilelang.testing from tilelang import tvm -from tilelang.layout import make_ascend_nz_layout +from tilelang.layout import make_ascend_l0c_layout from tvm import tirx from testing.ascend._ir import allocated_buffer, calls, nodes, statements @@ -25,7 +25,7 @@ def main(A: T.Tensor((32, 16), "bfloat16"), B: T.Tensor((n, 16), "bfloat16"), C: ub = T.view(storage, (2, 48, 80)) else: ub = T.alloc_shared((2, 48, 80), dtype) - T.annotate_layout({ub: make_ascend_nz_layout(ub)}) + T.annotate_layout({ub: make_ascend_l0c_layout(ub)}) T.copy(A, a) T.copy(B, b) T.gemm(a, b, acc, transpose_B=True, clear_accum=True) diff --git a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py index 8692dee50a..2e5584e20a 100644 --- a/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py +++ b/testing/ascend/layout/test_tilelang_ascend_fractal_layout_tags.py @@ -8,7 +8,6 @@ make_ascend_major_k_layout, make_ascend_major_mn_layout, make_ascend_sf_layout, - make_ascend_nz_layout, try_extract_fractal_layout, ) @@ -19,7 +18,6 @@ (make_ascend_major_k_layout, (16, 32), "float8_e4m3fn", (0, 1, 0)), (make_ascend_major_mn_layout, (32, 16), "float8_e4m3fn", (3, 0, 1)), (make_ascend_l0c_layout, (16, 16), "float32", (1, 1, 0)), - (make_ascend_nz_layout, (2, 32, 64), "float32", (1, 1, 0)), (make_ascend_sf_layout, (16, 1), "uint16", (2, 1, 0)), (make_ascend_sf_layout, (32, 1), "uint16", (2, 1, 0)), # The same small shape has different semantics despite identical indices. @@ -31,20 +29,6 @@ def test_canonical_layout_tags(factory, shape, dtype, expected): info = try_extract_fractal_layout(factory(buf), buf) assert (info.kind, info.c0_axis, info.row16_axis) == expected - if factory == make_ascend_nz_layout: - assert int(info.c0) == 16 - assert list(factory(buf).get_output_shape()) == [2, 4, 2, 16, 16] - - -@pytest.mark.parametrize("scope,kind", [("shared", 1), ("shared.dyn", 1), ("shared.l0c", 1), ("shared.l1", 0)]) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -def test_nz_layout_uses_buffer_scope(scope, kind, dtype): - buf = tirx.decl_buffer((2, 32, 64), dtype, name="A", scope=scope) - layout = make_ascend_nz_layout(buf) - info = try_extract_fractal_layout(layout, buf) - c0 = 8 if scope == "shared.l1" and dtype == "float32" else 16 - assert (info.kind, int(info.c0), info.c0_axis, info.row16_axis) == (kind, c0, 1, 0) - assert list(layout.get_output_shape()) == [2, 64 // c0, 2, 16, c0] def test_expand_preserves_sf_tag(): diff --git a/testing/ascend/transform/test_ascend_insert_nd2nz.py b/testing/ascend/transform/test_ascend_insert_nd2nz.py index 3e28614bd6..101173b966 100644 --- a/testing/ascend/transform/test_ascend_insert_nd2nz.py +++ b/testing/ascend/transform/test_ascend_insert_nd2nz.py @@ -3,7 +3,7 @@ import pytest from tilelang import tvm from tilelang.ascend import language as T, transform -from tilelang.layout import make_ascend_compact_nz_layout, make_ascend_nz_layout +from tilelang.layout import make_ascend_compact_nz_layout from tvm import tirx from testing.ascend._ir import calls @@ -66,26 +66,3 @@ def test_scatter_storage_requirements(kwargs, error): transform.InsertNd2Nz()(before) else: assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_scatter")) == 1 - - -@pytest.mark.parametrize("factory", [make_ascend_nz_layout, make_ascend_compact_nz_layout]) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -def test_post_copy_requires_compact_nz_panel_pitch(factory, dtype): - src = tirx.decl_buffer((33, 128), dtype, name="src", scope="shared.dyn") - dst = tirx.decl_buffer((32, 128), dtype, name="dst", scope="shared.l1") - region = tirx.BufferRegion(src, [tvm.ir.Range(0, 32), tvm.ir.Range(0, 128)]) - root = tirx.SBlock( - [], - [], - [], - "root", - tirx.Evaluate(T.copy(region, dst)), - alloc_buffers=[src, dst], - annotations={"layout_map": {src: factory(src), dst: make_ascend_nz_layout(dst)}}, - ) - before = tvm.IRModule({"main": tirx.PrimFunc([], tirx.SBlockRealize([], True, root))}) - if factory == make_ascend_nz_layout: - with pytest.raises(ValueError, match="requires compact NZ storage"): - transform.InsertNd2Nz()(before) - else: - assert len(calls(transform.InsertNd2Nz()(before), "tl.ascend_nd2nz_post_copy")) == 1 diff --git a/testing/ascend/transform/test_ascend_lower_copy.py b/testing/ascend/transform/test_ascend_lower_copy.py index faa0d7522c..206a488a24 100644 --- a/testing/ascend/transform/test_ascend_lower_copy.py +++ b/testing/ascend/transform/test_ascend_lower_copy.py @@ -167,7 +167,7 @@ def test_fix_nz_uses_physical_panel_pitch_and_slice_origin(dual, src_shape): src, dst, dst_region=_region(dst, [1, 32, 64], [1, 16, 16]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)}, annotations={"dual_dst_ctl": dual}, unit_flag_ctrl=3, ) @@ -185,14 +185,14 @@ def test_fix_nz_rejects_partial_column_panels(column, width): src, dst, dst_region=_region(dst, [32, width], [0, column]), - layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}, + layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)}, ) def test_fix_nz_rejects_l1_float32_panel_width(): src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared") - with pytest.raises(ValueError, match="use make_ascend_nz_layout"): + with pytest.raises(ValueError, match="use make_ascend_l0c_layout"): _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_major_k_layout(dst)}) @@ -200,7 +200,7 @@ def test_fix_nz_rejects_noncompact_logical_strides(): src = tirx.decl_buffer((32, 64), "float32", name="src", scope="shared.l0c") dst = tirx.decl_buffer((32, 64), "float32", name="dst", scope="shared", strides=(80, 1)) with pytest.raises(ValueError, match="requires dense logical strides"): - _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_nz_layout(dst)}) + _lower(src, dst, layouts={src: make_ascend_l0c_layout(src), dst: make_ascend_l0c_layout(dst)}) @pytest.mark.parametrize("scope", ["shared.l1", "shared.dyn"]) diff --git a/tilelang/layout/ascend.py b/tilelang/layout/ascend.py index 0316543e9e..6409f341b3 100644 --- a/tilelang/layout/ascend.py +++ b/tilelang/layout/ascend.py @@ -18,25 +18,16 @@ class _AscendFractalLayout(Layout): def make_ascend_nz_layout(buffer: BufferLikeType): - """Create the Ascend NZ layout appropriate for the buffer's scope. - - Logical ``[..., m, n]`` maps to - ``[..., n // C0, m // 16, m % 16, n % C0]``. L0C and UB - (``shared``/``shared.dyn``) use the fixed 16-column accumulator format. - Other scopes use the Cube input format, with - ``C0 = 32 bytes / element_size``. - - With ``T.annotate_layout({ub: make_ascend_nz_layout(ub)})``, L0C-to-UB - ``T.copy`` and ``T.dual_copy`` preserve NZ. Declare padding in the UB - shape; copy regions retain the allocation's physical panel pitch and - must cover whole 16-column panels. Buffer indices remain logical; - SIMD pointer increments access consecutive physical elements. + """ + Create an Ascend NZ (zN fractal) layout for a Cube operand buffer. - For compact UB packing before UB-to-L1 copies, use - :func:`make_ascend_compact_nz_layout` explicitly. + A logical ``[.., rows, cols]`` tile maps to physical + ``[.., rows / 16, cols / C0, 16, C0]`` -- a grid of 16xC0 fractals -- where + ``C0 = 32 bytes / element_size`` is the Cube unit's fractal width and the + forward map is ``(.., r, c) -> [.., r / 16, c / C0, r % 16, c % C0]``. Args: - buffer: BufferLikeType (rank >= 2). + buffer: BufferLikeType (rank >= 2, constant last-2 dims) Returns: Layout: The NZ fractal layout for the buffer. @@ -96,6 +87,11 @@ def make_ascend_l0c_layout(buffer: BufferLikeType): The physical shape is ``[.., N / 16, M / 16, 16, 16]`` and the forward map is ``(.., m, n) -> [.., n / 16, m / 16, m % 16, n % 16]``. + + This layout can also annotate a UB buffer to preserve NZ in L0C-to-UB + ``T.copy`` and ``T.dual_copy``. Declare padding in the buffer shape; + copy regions retain the physical panel pitch of that allocation and + must cover whole 16-column panels. Buffer indices remain logical. """ buf, _, _ = _get_buffer_info(buffer) return _ffi_api.make_ascend_l0c_layout(buf) From b1883197719eca22412d6e954728dbe134977a09 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Sun, 18 Oct 2026 05:06:54 +0800 Subject: [PATCH 07/12] [Ascend] Select AIV streams once before scheduled pipelines (cherry picked from commit 05120e8f6460449820e90a01a131d72c419c64fc) --- src/ascend/transform/lower_scheduled_tir.cc | 26 ++++++ .../ascend/auto_schedule/test_core_streams.py | 92 +++++++++++++++++++ 2 files changed, 118 insertions(+) create mode 100644 testing/ascend/auto_schedule/test_core_streams.py diff --git a/src/ascend/transform/lower_scheduled_tir.cc b/src/ascend/transform/lower_scheduled_tir.cc index a9a6aba8c7..5eb485fa71 100644 --- a/src/ascend/transform/lower_scheduled_tir.cc +++ b/src/ascend/transform/lower_scheduled_tir.cc @@ -51,6 +51,7 @@ #include "ascend/transform/attr.h" #include "op/builtin.h" #include "support/check.h" +#include "tir/transforms/ir_utils.h" #include "transform/common/attr.h" namespace tvm { @@ -690,7 +691,9 @@ Stmt LowerScheduledKernel(const Stmt &scheduled_body, Optional vector_body = LowerCore(scheduled_body, vector_mask, false, sid); Optional cube_body = LowerCore(scheduled_body, kCoreCube); + bool allow_shared_scalar_stream = false; if (!vector_body.has_value() && !cube_body.has_value()) { + allow_shared_scalar_stream = true; vector_body = LowerCore(scheduled_body, vector_mask, /*allow_shared_scalar_stream=*/true, sid); } @@ -698,6 +701,29 @@ Stmt LowerScheduledKernel(const Stmt &scheduled_body, UsesVar(vector_body.value(), [&](const VarNode *var) { return var == sid.get(); }); + if (needs_sid && vector_count == 2) { + // Select the subcore once, before entering its pipeline. Interleaving + // guarded AIV tasks would otherwise test sid throughout the issue stream. + // Extract each stream independently to retain its own scalar bindings and + // loop variables, then specialize the remaining shared sid expressions. + auto project = [&](CoreMask core) -> Optional { + Optional body = + LowerCore(scheduled_body, core, allow_shared_scalar_stream, sid); + if (body.has_value()) + return Substitute(body.value(), + {{sid, IntImm(sid.dtype(), VectorSid(core))}}); + return std::nullopt; + }; + Optional aiv0 = project(kCoreVector0); + Optional aiv1 = project(kCoreVector1); + ICHECK(aiv0.has_value() || aiv1.has_value()); + if (aiv0.has_value()) + vector_body = IfThenElse(sid == 0, aiv0.value(), aiv1); + else + vector_body = IfThenElse(sid == 1, aiv1.value()); + // Shared tasks can contain nested Bind/For/AllocBuffer definitions too. + vector_body = ConvertSSA(vector_body.value()); + } if (vector_body.has_value() && !external_sid.has_value() && (cube_body.has_value() || needs_sid)) { IterVar sid_thread(Range(0, vector_count), sid, IterVarType::kThreadIndex, diff --git a/testing/ascend/auto_schedule/test_core_streams.py b/testing/ascend/auto_schedule/test_core_streams.py new file mode 100644 index 0000000000..4786212bca --- /dev/null +++ b/testing/ascend/auto_schedule/test_core_streams.py @@ -0,0 +1,92 @@ +"""Lowering preserves each AIV's program with a single subcore selection.""" + +import pytest +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform +from tvm import tirx + +from testing.ascend._ir import calls, kernel, nodes, seq +from testing.ascend.auto_schedule._scheduled_ir import unit + + +@pytest.mark.parametrize("active", [(0, 1), (0,), (1,)]) +def test_subcore_selection_preserves_bindings_and_order(active): + sid = tirx.Var("sid", "int32") + route = tirx.Var("route", "int32") + i = tirx.Var("i", "int32") + tasks = [unit(tirx.Bind(route, sid + 10), core=5)] + for subcore in active: + for label in ("before", "after"): + tasks.append( + unit( + tirx.Evaluate(tirx.call_extern("int32", f"{label}_{subcore}", i, route, sid)), + core=1 if subcore == 0 else 4, + ) + ) + loop = tirx.For(i, 0, 4, tirx.ForKind.SERIAL, seq(*tasks)) + before = kernel(unit(loop, core=None)) + func = before["main"] + thread = tirx.IterVar(tvm.ir.Range(0, 2), sid, tirx.IterVar.ThreadIndex, "cthread") + before = tvm.IRModule({"main": func.with_body(tirx.AttrStmt(thread, "thread_extent", 2, func.body))}) + after = transform.LowerScheduledTIR()(before) + assert tirx.analysis.verify_ssa(after["main"]) + (vector,) = [block for block in nodes(after, tirx.SBlock) if block.name_hint == "VECTOR"] + (selection,) = [ + node for node in nodes(vector.body, tirx.IfThenElse) if any(var.same_as(sid) for var in nodes(node.condition, tirx.Var)) + ] + assert isinstance(selection.condition, tirx.EQ) and int(selection.condition.b) == 0 + analyzer = tvm.arith.Analyzer() + for subcore, branch in enumerate((selection.then_case, selection.else_case)): + assert branch is not None + (loop,) = nodes(branch, tirx.For) + # The selector encloses the whole loop. No per-task sid predicates or + # references to another AIV's scalar binding survive inside it. + assert not any(var.same_as(sid) for var in nodes(loop, tirx.Var)) + records = calls(loop, "tirx.call_extern") + assert [call.args[0].value for call in records] == ([f"before_{subcore}", f"after_{subcore}"] if subcore in active else []) + bindings = nodes(loop, tirx.Bind) + assert len(bindings) == 1 + assert int(analyzer.simplify(bindings[0].value)) == subcore + 10 + for call in records: + assert call.args[2].same_as(bindings[0].var) + assert int(call.args[3]) == subcore + + +def test_shared_simd_body_keeps_distinct_local_definitions(): + @T.prim_func + def main(A: T.Tensor((2, 128), "float32"), B: T.Tensor((2, 128), "float32")): + with T.MixedKernel(1) as (_, sid): + a = T.alloc_shared((128,), "float32") + b = T.alloc_shared((128,), "float32") + T.copy(A[sid, :], a) + with T.SimdVF(): + mask = T.simd.pset(32) + for i in range(2): + value = T.simd.vld(a[i * 64]) + doubled = T.simd.vadd(value, value, mask) + T.simd.vsts(b[i * 64], doubled, mask) + T.copy(b, B[sid, :]) + + snapshots = {} + + @tvm.ir.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + if info.name in ("tl.InsertSync", "tl.LowerScheduledTIR"): + snapshots[info.name] = mod + + # VF-local bindings and loops live inside one atomic scheduling task, so + # projecting only the outer schedule units must also freshen these locals. + with tvm.transform.PassContext(instruments=[Capture()]): + artifact = tilelang.lower(main, target="ascend") + for mod in snapshots.values(): + assert tirx.analysis.verify_ssa(mod["main"]) + after = snapshots["tl.LowerScheduledTIR"] + (vector,) = [block for block in nodes(after, tirx.SBlock) if block.name_hint == "VECTOR"] + (selection,) = nodes(vector.body, tirx.IfThenElse) + for branch in (selection.then_case, selection.else_case): + assert len(calls(branch, "tl.simd.vadd")) == 1 + assert len(nodes(branch, tirx.For)) == 1 + assert "__simd_vf__" in artifact.kernel_source From da2ed9172c0faa150431ca2bc2d5f5c1f4b74b73 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:52:50 +0800 Subject: [PATCH 08/12] [Ascend] Preserve manual pipe order across physical phase ties (cherry picked from commit 3c264f6f983787af9ba2f8336030cbbdc223b458) --- src/ascend/transform/auto_schedule.cc | 2 + .../auto_schedule/test_special_registers.py | 44 ++++++++++++++++--- tilelang/ascend/language/schedule_hint.py | 10 +++-- tilelang/ascend/transform/z3_scheduler.py | 6 ++- 4 files changed, 51 insertions(+), 11 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 1687d6b936..868f8ca40a 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -611,6 +611,7 @@ std::vector> ScheduleBuilder::Z3SchedulePython( resource_deps.emplace_back(i, j); } if (manual_schedule && + (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } @@ -1035,6 +1036,7 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, resource_deps.emplace_back(i, j); } if (manual_schedule && + (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 3fe14d22e5..01142a226c 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -1,5 +1,6 @@ """Special-register dependencies participate in scheduling and core placement.""" +import inspect from unittest.mock import patch import pytest @@ -160,21 +161,52 @@ def test_cube_mode_schedule_preserves_iteration_order(mode, enable_offset, start tvm.ir.assert_structural_equal(child.body.body, operations[index]) -def test_manual_stages_break_scalar_pipe_ties_by_iteration(): +@pytest.mark.parametrize("stages", [(0, 0), (0, 1), (6, 10), (1, 0)], ids=["same", "ascending", "nonadjacent", "descending"]) +def test_manual_scalar_issue_order_survives_stage_ties(stages): operations = [tirx.Evaluate(T.set_hf32_mode("nearest_even")), tirx.Evaluate(T.set_mmad_direction("n"))] - body = seq(*(unit(op, core=0, stage=stage, cost=(1, 1)) for stage, op in enumerate(operations))) + body = seq(*(unit(op, core=0, stage=stage, cost=(1, 1)) for stage, op in zip(stages, operations))) loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) before = kernel(unit(loop, core=None, stage=-1)) # These independent setters share PIPE_Scalar, which has no resource - # exclusion. The real solver chooses II=1, forcing equal issue phases. - # The earlier logical iteration takes precedence over per-pipe source order. + # exclusion. Ascending stages need distinct phases so the reconstruction + # tie-break cannot reverse manual source order; other stage orders may tie. after = transform.AutoSchedule()(before) children = after["main"].body.block.body.body.body.seq - assert [int(child.node["stage"]) for child in children] == [1, 0] - for child, op in zip(children, reversed(operations)): + assert [int(child.node["stage"]) for child in children] == list(stages) + for child, op in zip(children, operations): tvm.ir.assert_structural_equal(child.body.body, op) +@pytest.mark.parametrize("cores", [(1, 4), (1, 2), (2, 4)], ids=["aiv0-aiv1", "aiv0-aic", "aic-aiv1"]) +@pytest.mark.parametrize("in_loop", [False, True], ids=["sequence", "loop"]) +def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop): + sources = [tirx.decl_buffer((64,), "float32", name=f"A{i}") for i in range(2)] + buffers = [ + tirx.decl_buffer((64,), "float32", name=f"tile{i}", scope="shared.l1" if core == 2 else "shared.dyn") + for i, core in enumerate(cores) + ] + stages = (0, 1) if in_loop else (0, 0) + body = seq(*(unit(copy(source, buffer), core=core, stage=stage, cost=(1, 1)) + for source, buffer, core, stage in zip(sources, buffers, cores, stages))) + if in_loop: + loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) + body = unit(loop, core=None, stage=-1) + name = "z3_schedule_loop_python" if in_loop else "z3_schedule_seq_python" + original = getattr(z3_scheduler, name) + inputs = [] + + def capture(*args, **kwargs): + inputs.append(dict(inspect.signature(original).bind(*args, **kwargs).arguments)) + return original(*args, **kwargs) + + with patch.object(z3_scheduler, name, side_effect=capture): + transform.AutoSchedule()(kernel(body, buffers=buffers, params=sources)) + (actual,) = inputs + # Equal pipe names on separate cores do not impose manual issue order. + # Shared HBM resources may independently require resource exclusion. + assert actual["pipe_order_deps"] == [] + + @pytest.mark.parametrize("mode", ["hf32", "mmad-direction"]) def test_cube_mode_writers_follow_nested_gemm_core(mode): before, setter = _cube_mode_program(mode) diff --git a/tilelang/ascend/language/schedule_hint.py b/tilelang/ascend/language/schedule_hint.py index 51e0fdb371..4c15dfe662 100644 --- a/tilelang/ascend/language/schedule_hint.py +++ b/tilelang/ascend/language/schedule_hint.py @@ -122,10 +122,12 @@ def Stage(stage: int): ``T.Stage`` annotation, AutoSchedule treats that child list as manual. One scope may contain several consecutive statements; every scheduler task materialized from those statements receives the same stage. Unannotated - siblings default to stage 0. Within each hardware pipe, source order is the - manual issue order; tasks on different pipes remain freely reorderable, so - no separate order annotation is required. Nested child lists select - automatic or manual scheduling independently. + siblings default to stage 0. On each shared core and hardware pipe, source + order is the manual issue order. The scheduler separates issue phases when + a stage tie-break would reverse this order. Data and register dependencies + remain in force; an incompatible manual order fails scheduling. Tasks on + different pipes remain freely reorderable subject to those dependencies. + Nested child lists select automatic or manual scheduling independently. ``T.Stage`` must wrap complete scheduler tasks. It must not appear inside ``T.Task``, ``T.SimtVF``, ``T.SimdVF``, an SBlock, or a parallel loop. In diff --git a/tilelang/ascend/transform/z3_scheduler.py b/tilelang/ascend/transform/z3_scheduler.py index 51e9f02bf2..4f771a922b 100644 --- a/tilelang/ascend/transform/z3_scheduler.py +++ b/tilelang/ascend/transform/z3_scheduler.py @@ -427,7 +427,11 @@ def z3_schedule_loop_python( for i, stage in enumerate(manual_stages): solver.add(k_vars[i] == stage) for previous, current in pipe_order_deps: - solver.add(r_vars[previous] <= r_vars[current]) + # Reconstruction places larger stages first at equal physical + # phases. Separate ascending stages to preserve manual source + # order; this gap expresses ordering, not instruction latency. + gap = int(manual_schedule and enable_offset and manual_stages[previous] < manual_stages[current]) + solver.add(r_vars[previous] + gap <= r_vars[current]) buffer_vars = [z3.Int(f"buf_{i}", ctx) for i in range(len(buffer_sizes))] for i in range(len(buffer_sizes)): From 370f6ef97903a179eda61f7247291e14ebe26051 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 11:58:47 +0800 Subject: [PATCH 09/12] [Ascend] Format manual pipe-order regression (cherry picked from commit ca0b4ef54929fe565622ef8835b3277469814374) --- testing/ascend/auto_schedule/test_special_registers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 01142a226c..5795eeb820 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -186,8 +186,12 @@ def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop for i, core in enumerate(cores) ] stages = (0, 1) if in_loop else (0, 0) - body = seq(*(unit(copy(source, buffer), core=core, stage=stage, cost=(1, 1)) - for source, buffer, core, stage in zip(sources, buffers, cores, stages))) + body = seq( + *( + unit(copy(source, buffer), core=core, stage=stage, cost=(1, 1)) + for source, buffer, core, stage in zip(sources, buffers, cores, stages) + ) + ) if in_loop: loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) body = unit(loop, core=None, stage=-1) From 081b4db43f96b71553353aa4684778e41bc16956 Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:04:11 +0800 Subject: [PATCH 10/12] [Ascend] Reuse the existing physical core and pipe masks (cherry picked from commit da7716998c36761c2845377c1de2e3e6b20eac8a) --- src/ascend/transform/auto_schedule.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ascend/transform/auto_schedule.cc b/src/ascend/transform/auto_schedule.cc index 868f8ca40a..1687d6b936 100644 --- a/src/ascend/transform/auto_schedule.cc +++ b/src/ascend/transform/auto_schedule.cc @@ -611,7 +611,6 @@ std::vector> ScheduleBuilder::Z3SchedulePython( resource_deps.emplace_back(i, j); } if (manual_schedule && - (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } @@ -1036,7 +1035,6 @@ void ScheduleBuilder::Z3SchedulePythonLoop(ControlNode *ctrl, resource_deps.emplace_back(i, j); } if (manual_schedule && - (NodeCoreMask(nodes[i]) & NodeCoreMask(nodes[j])) != 0 && (NodeResourceMask(nodes[i]) & NodeResourceMask(nodes[j])) != 0) { pipe_order_deps.emplace_back(i, j); } From a23a56494774b75495027c87f1e12203aa86727e Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Wed, 21 Oct 2026 12:05:47 +0800 Subject: [PATCH 11/12] [Ascend] Capture the sequence solver in core-order coverage (cherry picked from commit 9ec3fc08a8b5d3884621e71ee4ef3885c7b2236f) --- testing/ascend/auto_schedule/test_special_registers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testing/ascend/auto_schedule/test_special_registers.py b/testing/ascend/auto_schedule/test_special_registers.py index 5795eeb820..7342a0dc84 100644 --- a/testing/ascend/auto_schedule/test_special_registers.py +++ b/testing/ascend/auto_schedule/test_special_registers.py @@ -195,7 +195,7 @@ def test_manual_pipe_order_does_not_order_distinct_physical_cores(cores, in_loop if in_loop: loop = tirx.For(tirx.Var("i", "int32"), 0, 3, tirx.ForKind.SERIAL, body, annotations={"enable_offset": True}) body = unit(loop, core=None, stage=-1) - name = "z3_schedule_loop_python" if in_loop else "z3_schedule_seq_python" + name = "z3_schedule_loop_python" if in_loop else "z3_schedule_python" original = getattr(z3_scheduler, name) inputs = [] From 9609e7ad512e104c24418ef303ed618c31be247d Mon Sep 17 00:00:00 2001 From: Denver Jin Date: Thu, 22 Oct 2026 17:03:52 +0800 Subject: [PATCH 12/12] [Ascend] Update FlashMLA compiler dependency status --- examples/ascend/flashmla/README.md | 7 +++--- examples/ascend/flashmla/WORKTREES.md | 33 +++++++++++++++------------ 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/examples/ascend/flashmla/README.md b/examples/ascend/flashmla/README.md index b6cca7cf3a..e741fed6da 100644 --- a/examples/ascend/flashmla/README.md +++ b/examples/ascend/flashmla/README.md @@ -66,9 +66,10 @@ fixed loop extent; they do not shorten this example's pipeline loop. ## Reproduce -This example is proposed separately from its compiler prerequisites, listed in -[WORKTREES.md](WORKTREES.md). The example branch starts at main `d85a70a`; -apply the prerequisite compiler changes before compiling the kernels. +The kernel branch includes the native NZ-copy, AIV-stream-selection, and manual +pipe-order compiler changes. It starts at main `d85a70a`; apply the remaining +compiler prerequisites and placement-API migration listed in +[WORKTREES.md](WORKTREES.md) before compiling the staged kernel. The historical integrated compiler was based on main `876a59e`; its [rebase validation record](results/staged-main-rebase-2026-10-20.json) covers the initial native build, host tests, and kernel source generation. The diff --git a/examples/ascend/flashmla/WORKTREES.md b/examples/ascend/flashmla/WORKTREES.md index e7134aaefd..38bd906dd4 100644 --- a/examples/ascend/flashmla/WORKTREES.md +++ b/examples/ascend/flashmla/WORKTREES.md @@ -1,31 +1,36 @@ # FlashMLA compiler prerequisites -The example branch contains the kernels, runners, and historical measurement -records. Compiler changes are proposed separately against `Denverjin/tilelang:main`, -synchronized to upstream `d85a70a97218fec25856f8b5c969f167ad96dc1e` on 2026-10-21. -Apply the compiler prerequisites before building the staged kernel. +The kernel branch contains the kernels, runners, historical measurement records, +and the compiler changes from fork PRs #5, #7, and #12. It starts from +`d85a70a97218fec25856f8b5c969f167ad96dc1e`, synchronized with upstream on 2026-10-21. +Apply the remaining compiler prerequisites before building the staged kernel. The source does not use `T.IssuePhase`, and its API is excluded from these proposals. -| Compiler change | Pull request | Fork branch | +| Included compiler change | Source pull request | +| --- | --- | +| Support native NZ output for L0C-to-UB copies | [Denverjin/tilelang#5](https://github.com/Denverjin/tilelang/pull/5) | +| Select AIV streams outside scheduled pipelines | [Denverjin/tilelang#7](https://github.com/Denverjin/tilelang/pull/7) | +| Preserve manual source order across physical phase ties | [Denverjin/tilelang#12](https://github.com/Denverjin/tilelang/pull/12) | + +| Remaining compiler change | Pull request | Fork branch | | --- | --- | --- | -| Support native NZ output for L0C-to-UB copies | [#5](https://github.com/Denverjin/tilelang/pull/5) | [codex/ascend-l0c-ub-nz](https://github.com/Denverjin/tilelang/tree/codex/ascend-l0c-ub-nz) | | Preserve Cube configuration phases across manual stages | [#6](https://github.com/Denverjin/tilelang/pull/6) | [codex/ascend-cube-register-phases](https://github.com/Denverjin/tilelang/tree/codex/ascend-cube-register-phases) | -| Select AIV streams outside scheduled pipelines | [#7](https://github.com/Denverjin/tilelang/pull/7) | [codex/ascend-aiv-stream-selection](https://github.com/Denverjin/tilelang/tree/codex/ascend-aiv-stream-selection) | -| Preserve complete periodic buffer lifetimes | [#8](https://github.com/Denverjin/tilelang/pull/8) | [codex/ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/codex/ascend-periodic-buffer-lifetimes) | -| Establish buffer placement hints from exact L1 offsets | [#9](https://github.com/Denverjin/tilelang/pull/9) | [codex/ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/codex/ascend-buffer-placement-hints) | +| Preserve complete periodic buffer lifetimes | [tile-ai/tilelang#3606](https://github.com/tile-ai/tilelang/pull/3606), from [fork #8](https://github.com/Denverjin/tilelang/pull/8) | [ascend-periodic-buffer-lifetimes](https://github.com/Denverjin/tilelang/tree/ascend-periodic-buffer-lifetimes) | +| Support exact placement of versioned on-chip buffers | [tile-ai/tilelang#3607](https://github.com/tile-ai/tilelang/pull/3607), from [fork #9](https://github.com/Denverjin/tilelang/pull/9) | [ascend-buffer-placement-hints](https://github.com/Denverjin/tilelang/tree/ascend-buffer-placement-hints) | | Preserve physical views and pointer units during layout lowering | [#10](https://github.com/Denverjin/tilelang/pull/10) | [codex/ascend-layout-pointer-units](https://github.com/Denverjin/tilelang/tree/codex/ascend-layout-pointer-units) | | Preserve independent Cube view layouts on shared storage rings | [#11](https://github.com/Denverjin/tilelang/pull/11) | [codex/ascend-conflict-storage-rings](https://github.com/Denverjin/tilelang/tree/codex/ascend-conflict-storage-rings) | -| Preserve manual source order across physical phase ties | [#12](https://github.com/Denverjin/tilelang/pull/12) | [codex/ascend-manual-pipe-order](https://github.com/Denverjin/tilelang/tree/codex/ascend-manual-pipe-order) | The additional NZ pitch, compact-allocation, and access-footprint prerequisite is already proposed upstream as [tile-ai/tilelang#3530](https://github.com/tile-ai/tilelang/pull/3530). It is not duplicated as a new fork PR. The pointer-unit branch also carries physical allocation/view remapping fixes that are outside that upstream PR. -The buffer-placement branch currently implements exact L1 byte offsets. -Its PR describes the intended extension to manual placement hints for every -buffer kind. That extension is not implemented yet; the kernel still calls -`T.annotate_l1_offsets` and uses the existing L1 placement contract. +The buffer-placement prerequisite now provides `T.annotate_buffer_offsets` for +UB, L1, L0A/L0B/L0C, and SS buffers, with an explicit fixed +`T.annotate_buffer_versions` count required for every placed buffer. The staged +kernel still calls the earlier `T.annotate_l1_offsets` API. Integrating the new +placement prerequisite also requires migrating that call and adding the explicit +fixed version count for `q_l1` on the prefetched-query path. The staged kernel requires PR #11's Cube-view layout support from commit `393a0bde46d60a3b7cfaabeeb65ebc67a7a14ec7` or later. QK allocates one