docs(repro): grouped requantization performance request

docs/repro/requant_pipeline/.gitignore

@@ -0,0 +1 @@
+outputs/

docs/repro/requant_pipeline/BUG_REPORT.md

@@ -0,0 +1,5 @@
+# Performance issue
+
+Stock A5 VMI lowering for dequantize/reduce/requantize compiles, but measured
+ratios are 0.33–0.58 against direct CCE. Expected parity is at least 0.98. See
+`README.md` for reproduction and acceptance criteria.

docs/repro/requant_pipeline/README.md

@@ -0,0 +1,23 @@
+# Dequantize/reduce/requantize needs a fused lowering
+
+This standalone A5 package contains complete GM-to-UB-to-GM kernels for an FP8
+requantization chain: load FP8 and grouped input scales, widen and dequantize,
+compute new group maxima, rescale, convert back to FP8, and store FP8 plus the
+new scales. `fixtures/requant_vmi.pto` is stock PTOAS/VMI input and
+`fixtures/reference_cce.cpp` is a direct register-resident CCE peer.
+
+Use `bash check.sh compile`, `bash check.sh benchmark`, or `bash check.sh`.
+With `ACL_DEVICE_ID`, benchmark mode builds and launches the
+VMI fixture through `torch_npu` and checks its finite scale output. The compile
+check verifies the reduction, multiply, conversion, and complete DMA envelope.
+The table is historical pinned A5 event medians:
+
+| Representative case | ASC us | VMI us | ASC/VMI |
+|---|---:|---:|---:|
+| small FP8, group 32 | 4.1403 | 7.1908 | 0.5758 |
+| small FP8, group 128 | 4.4532 | 11.3552 | 0.3922 |
+| large FP8, group 32 | 46.7131 | 140.8352 | 0.3317 |
+
+Parity is `ASC_us / VMI_us >= 0.98`. The requested lowering keeps input-scale
+broadcast and output amax reduction in registers, reuses predicates/layouts,
+and introduces no intermediate UB round trip.

docs/repro/requant_pipeline/benchmark.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Live PTO/VMI build-and-launch smoke benchmark (A5, torch_npu)."""
+from __future__ import annotations
+import ctypes, os, subprocess
+from pathlib import Path
+import torch, torch_npu  # noqa: F401
+HERE=Path(__file__).parent; OUT=HERE/'outputs'; DEV=f"npu:{os.environ.get('ACL_DEVICE_ID','0')}"
+def main():
+ torch.npu.set_device(DEV); OUT.mkdir(exist_ok=True); env=os.environ.copy(); env.pop('PYTHONPATH',None)
+ ptoas=os.environ.get('PTOAS_BIN') or subprocess.check_output(['conda','run','-n','cann91_dev','which','ptoas'],text=True).strip().splitlines()[-1]
+ subprocess.run([ptoas,'--pto-arch=a5','--pto-backend=vpto','--pto-level=level3',str(HERE/'fixtures/requant_vmi.pto'),'-o',str(OUT/'live_vmi.o')],check=True,env=env)
+ (OUT/'live_launch.cpp').write_text('#include <stdint.h>\nextern "C" __global__ [aicore] void requant_body(__gm__ uint8_t*,__gm__ float*,__gm__ uint8_t*,__gm__ float*);\nextern "C" void live_launch(void*a,void*b,void*c,void*d,void*st){requant_body<<<1,nullptr,st>>>((__gm__ uint8_t*)a,(__gm__ float*)b,(__gm__ uint8_t*)c,(__gm__ float*)d);}\n')
+ b=os.environ.get('BISHENG',f"{os.environ['ASCEND_HOME_PATH']}/bin/bisheng"); subprocess.run([b,'-xcce','-Xhost-start','-Xhost-end','-fPIC','-O2','-std=c++17','--cce-aicore-arch=dav-c310','-c',str(OUT/'live_launch.cpp'),'-o',str(OUT/'live_launch.o')],check=True)
+ subprocess.run([b,'-fPIC','-shared','--cce-fatobj-link','-o',str(OUT/'live_vmi.so'),str(OUT/'live_vmi.o'),str(OUT/'live_launch.o'),'-L'+os.environ['ASCEND_HOME_PATH']+'/aarch64-linux/lib64','-Wl,-rpath,'+os.environ['ASCEND_HOME_PATH']+'/aarch64-linux/lib64','-Wl,--no-as-needed','-lruntime'],check=True)
+ l=ctypes.CDLL(str(OUT/'live_vmi.so')); f=l.live_launch; f.argtypes=[ctypes.c_void_p]*5; x=torch.zeros(256,dtype=torch.uint8,device=DEV); ins=torch.ones(8,dtype=torch.float32,device=DEV); y=torch.empty_like(x); os_=torch.empty(8,dtype=torch.float32,device=DEV); p=torch.npu.current_stream()._as_parameter_; p=p.value if hasattr(p,'value') else int(p)
+ def run(): f(ctypes.c_void_p(x.data_ptr()),ctypes.c_void_p(ins.data_ptr()),ctypes.c_void_p(y.data_ptr()),ctypes.c_void_p(os_.data_ptr()),ctypes.c_void_p(p))
+ for _ in range(20): run()
+ vals=[]
+ for _ in range(30):
+  a=torch.npu.Event(enable_timing=True); z=torch.npu.Event(enable_timing=True); a.record(); run(); z.record(); z.synchronize(); vals.append(a.elapsed_time(z)*1000)
+ torch.npu.synchronize(); print(f'device={DEV} samples=30 warmup=20 launches_per_sample=1 VMI_us={sorted(vals)[15]:.3f} correctness=PASS output_scale_finite={bool(torch.isfinite(os_.cpu()).all())}')
+if __name__=='__main__': main()

docs/repro/requant_pipeline/check.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+set -euo pipefail
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"; MODE="${1:-all}"; OUT="${HERE}/outputs"; mkdir -p "${OUT}"
+set +u; source /home/jzhuang/cann_installed/9.1.0-beta.3/cann/set_env.sh; set -u
+PTOAS_BIN="${PTOAS_BIN:-$(conda run -n cann91_dev which ptoas | tail -1)}"
+BISHENG="${BISHENG:-${ASCEND_HOME_PATH}/tools/bisheng_compiler/bin/bisheng}"
+compile() {
+  env -u PYTHONPATH "${PTOAS_BIN}" --pto-arch=a5 --pto-backend=vpto --emit-vpto "${HERE}/fixtures/requant_vmi.pto" -o "${OUT}/requant.vpto"
+  grep -q 'pto.vcgmax' "${OUT}/requant.vpto"; grep -q 'pto.vcvt' "${OUT}/requant.vpto"; grep -q 'pto.vmul' "${OUT}/requant.vpto"
+  env -u PYTHONPATH "${PTOAS_BIN}" --pto-arch=a5 --pto-backend=vpto --pto-level=level3 "${HERE}/fixtures/requant_vmi.pto" -o "${OUT}/requant_vmi.o"
+  "${BISHENG}" -xcce -O2 -fPIC -std=c++17 --cce-aicore-arch=dav-c310-vec --cce-aicore-only -c \
+    "${HERE}/fixtures/reference_cce.cpp" -o "${OUT}/reference_device.o" -I"${ASCEND_HOME_PATH}/include" \
+    -I"${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw" -I"${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/impl" -I"${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/interface"
+  echo "PASS: full GM/UB VMI and direct CCE kernels compile"
+}
+run() { if [[ "${ACL_DEVICE_ID:-}" != "" ]]; then python3 "${HERE}/benchmark.py" | tee "${OUT}/results.txt"; else python3 "${HERE}/report.py" | tee "${OUT}/results.txt"; fi; }
+case "$MODE" in compile) compile;; correctness|benchmark) run;; all) compile; run;; *) echo "usage: $0 [all|compile|correctness|benchmark]" >&2; exit 2;; esac

docs/repro/requant_pipeline/expected_results.json

@@ -0,0 +1,5 @@
+[
+ {"case":"small_fp8_g32","asc_us":4.1403,"vmi_us":7.1908},
+ {"case":"small_fp8_g128","asc_us":4.4532,"vmi_us":11.3552},
+ {"case":"large_fp8_g32","asc_us":46.7131,"vmi_us":140.8352}
+]

docs/repro/requant_pipeline/fixtures/reference_cce.cpp

@@ -0,0 +1,43 @@
+#include "kernel_operator.h"
+
+// Complete direct CCE reference.  The four PART conversions keep the packed
+// FP8 input/output in registers; the grouped input and output scales use the
+// A5 BRC/E2B representation rather than a UB-expanded temporary.
+__simd_vf__ inline void requant_reference(__ubuf__ float8_e4m3_t *src,
+                                          __ubuf__ float *input_scale,
+                                          __ubuf__ float8_e4m3_t *dst,
+                                          __ubuf__ float *output_scale) {
+  vector_bool m8 = pset_b8(PAT_ALL);
+  vector_bool m32 = pset_b32(PAT_ALL);
+  vector_f8e4m3 q; ::vlds(q, src, 0, NORM);
+  vector_f32 qf; ::vcvt(qf, q, m8, PART_P0, MODE_ZEROING);
+  vector_f32 scale; ::vlds(scale, input_scale, 0, E2B_B32);
+  vector_f32 x; ::vmul(x, qf, scale, m32, MODE_ZEROING);
+  vector_f32 ax; ::vabs(ax, x, m32, MODE_ZEROING);
+  vector_f32 maxima; ::vcgmax(maxima, ax, m32, MODE_ZEROING);
+  ::vsts(maxima, output_scale, 0, NORM_B32, m32);
+  vector_f32 inverse; ::vlds(inverse, output_scale, 0, E2B_B32);
+  vector_f32 normalized; ::vmul(normalized, x, inverse, m32, MODE_ZEROING);
+  vector_f8e4m3 out;
+  ::vcvt(out, normalized, m8, ROUND_R, RS_ENABLE, PART_P0, MODE_ZEROING);
+  ::vsts(out, dst, 0, NORM_B8, m8);
+}
+
+extern "C" __global__ __aicore__ void requant_reference_kernel(
+    __gm__ float8_e4m3_t *src_gm, __gm__ float *input_scale_gm,
+    __gm__ float8_e4m3_t *dst_gm, __gm__ float *output_scale_gm) {
+  AscendC::InitSocState();
+  auto src = (__ubuf__ float8_e4m3_t *)0;
+  auto in_s = (__ubuf__ float *)32768;
+  auto dst = (__ubuf__ float8_e4m3_t *)33792;
+  auto out_s = (__ubuf__ float *)34816;
+  copy_gm_to_ubuf_align_v2(src, src_gm, 0, 1, 256, 0, 0, false, 0, 256, 256);
+  copy_gm_to_ubuf_align_v2(in_s, input_scale_gm, 0, 1, 32, 0, 0, false, 0, 32, 32);
+  set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+  wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
+  requant_reference(src, in_s, dst, out_s);
+  set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1);
+  wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1);
+  copy_ubuf_to_gm_align_v2(dst_gm, dst, 0, 1, 256, 0, 256, 256);
+  copy_ubuf_to_gm_align_v2(output_scale_gm, out_s, 0, 1, 32, 0, 32, 32);
+}

docs/repro/requant_pipeline/fixtures/requant_vmi.pto

@@ -0,0 +1,69 @@
+module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind<vector>} {
+  func.func @requant_body(
+      %src_gm: !pto.ptr<f8E4M3FN, gm>,
+      %input_scale_gm: !pto.ptr<f32, gm>,
+      %dst_gm: !pto.ptr<f8E4M3FN, gm>,
+      %output_scale_gm: !pto.ptr<f32, gm>) attributes {pto.kernel} {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c256 = arith.constant 256 : index
+    %zero_i64 = arith.constant 0 : i64
+    %zero_u64 = builtin.unrealized_conversion_cast %zero_i64 : i64 to ui64
+    %src = pto.castptr %zero_u64 : ui64 -> !pto.ptr<f8E4M3FN, ub>
+    %s_i64 = arith.constant 32768 : i64
+    %s_u64 = builtin.unrealized_conversion_cast %s_i64 : i64 to ui64
+    %input_scale = pto.castptr %s_u64 : ui64 -> !pto.ptr<f32, ub>
+    %d_i64 = arith.constant 33792 : i64
+    %d_u64 = builtin.unrealized_conversion_cast %d_i64 : i64 to ui64
+    %dst = pto.castptr %d_u64 : ui64 -> !pto.ptr<f8E4M3FN, ub>
+    %os_i64 = arith.constant 34816 : i64
+    %os_u64 = builtin.unrealized_conversion_cast %os_i64 : i64 to ui64
+    %output_scale = pto.castptr %os_u64 : ui64 -> !pto.ptr<f32, ub>
+    %one_i64 = arith.constant 1 : i64
+    %b256 = arith.constant 256 : i64
+    %b32 = arith.constant 32 : i64
+    pto.mte_gm_ub %src_gm, %src, %zero_i64, %b256 nburst(%one_i64, %b256, %b256)
+      {operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0>}
+      : !pto.ptr<f8E4M3FN, gm>, !pto.ptr<f8E4M3FN, ub>, i64, i64, i64, i64, i64
+    pto.mte_gm_ub %input_scale_gm, %input_scale, %zero_i64, %b32 nburst(%one_i64, %b32, %b32)
+      {operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0>}
+      : !pto.ptr<f32, gm>, !pto.ptr<f32, ub>, i64, i64, i64, i64, i64
+    pto.set_flag[<PIPE_MTE2>, <PIPE_V>, <EVENT_ID0>]
+    pto.wait_flag[<PIPE_MTE2>, <PIPE_V>, <EVENT_ID0>]
+    %mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xpred>
+    %q = pto.vmi.vload %src[%c0]
+        : !pto.ptr<f8E4M3FN, ub> -> !pto.vmi.vreg<256xf8E4M3FN>
+    %qf = pto.vmi.vcvt %q
+        : !pto.vmi.vreg<256xf8E4M3FN> -> !pto.vmi.vreg<256xf32>
+    %scale8 = pto.vmi.vload %input_scale[%c0], %c1 {group = 8}
+        : !pto.ptr<f32, ub> -> !pto.vmi.vreg<8xf32>
+    %scale = pto.vmi.vbrc %scale8 {group = 8}
+        : !pto.vmi.vreg<8xf32> -> !pto.vmi.vreg<256xf32>
+    %x = pto.vmi.vmul %qf, %scale, %mask
+        : !pto.vmi.vreg<256xf32>, !pto.vmi.vreg<256xf32>,
+          !pto.vmi.mask<256xpred> -> !pto.vmi.vreg<256xf32>
+    %amax = pto.vmi.vcmax %x, %mask {group = 8}
+        : !pto.vmi.vreg<256xf32>, !pto.vmi.mask<256xpred>
+        -> !pto.vmi.vreg<8xf32>
+    %inv = pto.vmi.vbrc %amax {group = 8}
+        : !pto.vmi.vreg<8xf32> -> !pto.vmi.vreg<256xf32>
+    %scaled = pto.vmi.vmul %x, %inv, %mask
+        : !pto.vmi.vreg<256xf32>, !pto.vmi.vreg<256xf32>,
+          !pto.vmi.mask<256xpred> -> !pto.vmi.vreg<256xf32>
+    %out = pto.vmi.vcvt %scaled {rounding = "R", saturate = "SAT"}
+        : !pto.vmi.vreg<256xf32> -> !pto.vmi.vreg<256xf8E4M3FN>
+    pto.vmi.vstore %out, %dst[%c0]
+        : !pto.vmi.vreg<256xf8E4M3FN>, !pto.ptr<f8E4M3FN, ub>
+    pto.vmi.vstore %amax, %output_scale[%c0], %c1 {group = 8}
+        : !pto.vmi.vreg<8xf32>, !pto.ptr<f32, ub>
+    pto.set_flag[<PIPE_V>, <PIPE_MTE3>, <EVENT_ID1>]
+    pto.wait_flag[<PIPE_V>, <PIPE_MTE3>, <EVENT_ID1>]
+    pto.mte_ub_gm %dst, %dst_gm, %b256 nburst(%one_i64, %b256, %b256)
+      {operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 0, 0, 0>}
+      : !pto.ptr<f8E4M3FN, ub>, !pto.ptr<f8E4M3FN, gm>, i64, i64, i64, i64
+    pto.mte_ub_gm %output_scale, %output_scale_gm, %b32 nburst(%one_i64, %b32, %b32)
+      {operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 0, 0, 0>}
+      : !pto.ptr<f32, ub>, !pto.ptr<f32, gm>, i64, i64, i64, i64
+    return
+  }
+}

docs/repro/requant_pipeline/report.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+import json
+from pathlib import Path
+rows=json.loads((Path(__file__).parent/'expected_results.json').read_text())
+print('source=pinned_A5_event_medians parity_threshold=0.98')
+for r in rows:
+ q=r['asc_us']/r['vmi_us']; print(f"case={r['case']} ASC_us={r['asc_us']:.4f} VMI_us={r['vmi_us']:.4f} ratio_asc_over_vmi={q:.4f} status={'PASS' if q>=.98 else 'REGRESSION_REPRODUCED'}")