[Bug][simpler][a5sim] Simulated kernel workers do not select the PTO ISA A5 architecture

[Crystal-wzy]

Body

Problem

Selecting platform="a5sim" leaves the kernel worker's PTO ISA CPU memory
model at its default, NPUArch::A2A3. As a result, TROWSUM executes the
A2A3 reduction path even though the selected platform is A5.

Expected: each worker selects NPUArch::A5 before using ISA Tiles.
Observed: a fixed FP32 row-sum input returns 1 instead of 0. Initializing
the worker as A5 fixes the result without changing the reduction or ISA headers.
In attention models, this difference can propagate through BF16 rounding
to cache and compressor-state outputs.

Environment

Component	Revision / version
pypto-lib	d6f2920046df7a19c28a5079c096a53b7e600219
PyPTO	8ecc5d19f0ac1062ae6779ba8aca8cee75c22667
simpler	39ce891dbb3f665e72e99b4a1387b47012fb18bd + PR #2186 patches listed below
pto-isa	bc1bd8b9544502a352c6cec1776b9ef00a8303f4
PTOAS	v0.61
PyTorch	2.6.0+cpu
Python	3.10.19
CPU compiler	GCC 15.2.1, C++23, -O2, without -ffast-math

The two simpler PR #2186 patches are
14a46178d83ad31c69ae06adacad2113325206a0 and
a9a149a4033cd49788582e9f109167c148b00579.
These are CPU simulator tests; no NPU or model weights are required.

Error case

From the pypto-lib repository root, with the dependencies above installed:

export PYTHONPATH=".${PYTHONPATH:+:$PYTHONPATH}"
export DEEPSEEK_V4_VARIANT=pro
timeout --kill-after=10s 300s python \
  models/deepseek_v4_pro/decode_attention_csa.py -p a5sim
timeout --kill-after=10s 300s python \
  models/deepseek_v4_pro/decode_attention_swa.py -p a5sim

Model	Outputs affected by this issue
decode_attention_csa.py	compress_state, inner_compress_state, cmp_kv
decode_attention_swa.py	kv_cache

These entrypoints generate random data, so they may pass on some runs.
Model controls below replayed captured failing inputs with the same golden
data and tolerances. The minimal reproduction uses fixed data and needs
no captured files. SWA's separate x_out failure is covered by
subissue_pypto_missing_gm_cross_core_sync.md.

Minimal reproduction

Save this as repro_simpler_a5_arch.py in the pypto-lib repository root.
PyPTO generates a single row-sum kernel and launches it through simpler;
golden is pypto-lib's validation helper.

Each of the eight input rows contains [1e8, 1, -1e8, 1] repeated 16 times.
For this input, A5's modeled FP32 adjacent-pair reduction yields zero; the
observed A2A3 path yields one with the listed compiler flags. Eight rows
satisfy the generated reduction tile's 32-byte column alignment requirement.

The optional --initialize-a5 control inserts A5 initialization at the
generated kernel entry, before Tile binding. It is a diagnostic hook for
the listed simpler version; it leaves the input, reduction, golden and
tolerances unchanged.

"""Reproduce a5sim failing to select the PTO ISA A5 reduction path."""
import argparse
from pathlib import Path
import re

import torch
import pypto.language as pl
from golden import TensorSpec, run


@pl.jit
def row_sum(x: pl.Tensor[[8, 64], pl.FP32],
            out: pl.Out[pl.Tensor[[8, 1], pl.FP32]]):
    with pl.spmd(1):
        idx = pl.tile.get_block_idx()
        out[:, :] = pl.row_sum(x[:, :])
    return out


def golden(tensors):
    # Expected from the A5 FP32 adjacent-pair reduction for this exact input.
    tensors["out"].zero_()


def initialize_a5_for_diagnosis():
    from simpler_setup.kernel_compiler import KernelCompiler

    original_compile = KernelCompiler._compile_incore_sim

    def compile_with_a5(self, source_path, *args, **kwargs):
        path = Path(source_path)
        source, count = re.subn(
            r"(void kernel_entry\([^\n]*\)\s*\{)",
            r"\1\n    pto::NPUMemoryModel::Instance().Initialize(pto::NPUArch::A5);",
            path.read_text(),
        )
        assert count == 1, "Expected one generated kernel entry"
        diagnostic = path.with_name(path.stem + ".a5-arch.cpp")
        diagnostic.write_text(source)
        return original_compile(self, str(diagnostic), *args, **kwargs)

    KernelCompiler._compile_incore_sim = compile_with_a5


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--initialize-a5", action="store_true")
    args = parser.parse_args()
    if args.initialize_a5:
        initialize_a5_for_diagnosis()
    result = run(
        fn=row_sum,
        specs=[
            TensorSpec("x", [8, 64], torch.float32,
                       init_value=lambda: torch.tensor([1e8, 1.0, -1e8, 1.0]).repeat(8, 16)),
            TensorSpec("out", [8, 1], torch.float32),
        ],
        golden_fn=golden,
        config={"platform": "a5sim"},
        atol=0,
        rtol=0,
    )
    if result.error:
        print(result.error)
    print(f"passed={result.passed}")
    raise SystemExit(0 if result.passed else 1)

Run the two modes in separate processes. If an external compilation cache
is enabled, use fresh outputs so both modes compile their kernel entry.

export PYTHONPATH=".${PYTHONPATH:+:$PYTHONPATH}"
python repro_simpler_a5_arch.py
echo "rc=$?"
python repro_simpler_a5_arch.py --initialize-a5
echo "rc=$?"

Both modes use atol=rtol=0. Recorded results:

Mode	Expected output	Actual output	Validation	Exit code
Default a5sim	Eight zeros	Eight ones	8/8 mismatches	1
a5sim --initialize-a5	Eight zeros	Eight zeros	PASS	0

Root cause and ownership

The failure follows this initialization path:

1. In simpler, Gxx15Toolchain (simpler_setup/toolchain.py) sets CPU_SIM
   and AIC/AIV compilation flags. KernelCompiler._compile_incore_sim
   (simpler_setup/kernel_compiler.py) and the tested simulator runtime
   do not select the A5 memory model for the executing worker.
2. In pto-isa, include/pto/cpu/NPUMemoryModel.hpp provides thread-local
   instances and defaults to NPUArch::A2A3 when none is explicitly selected.
3. include/pto/cpu/TRowSum.hpp chooses its reduction path using that
   instance's GetArch(). Consequently, this A5 launch uses the A2A3 path.

The A5 reduction algorithm was already fixed in pto-isa
9aba183fbb76f758272453ea5341e0c6185fa169.
The remaining integration must enable it. The ISA contract in
docs/coding/cpu_sim.md, under "Selecting the simulated architecture",
requires external integrations to select the target before running kernels.
Together with the initialization-only control, this identifies simpler's
worker initialization as the fix location.

Proposed fix

Map a5sim to NPUArch::A5 and a2a3sim to NPUArch::A2A3, then initialize
the ISA instance used by each kernel worker before its first Tile binding.
A kernel-entry adapter or a worker-invoked initializer in the kernel's
shared library can provide this. The reproduction hook validates this
direction; it is not the production implementation.

The implementation should:

• Cover generated and externally supplied kernels, worker reuse and
  separately loaded kernel shared libraries. Initializing only the main
  thread does not initialize worker-local instances.
• Preserve target selection for both platforms and initialize before any
  live Tile references the memory being configured.
• Update compilation cache identity when wrappers or flags change.

Verification and remaining scope

The standalone control above passes exactly after A5 initialization.
Model controls confirm its effect on the original failures:

Fixed input	Before initialization	After initialization
CSA, captured with seed 1	RMSNorm differs from its reference in 68/57344 BF16 elements; maximum absolute difference 0.0078125	All 68 differences disappear; all 7 model outputs pass
SWA attention, captured with seed 2	kv_cache FAIL, x_out FAIL	kv_cache PASS, x_out still FAIL from the independent synchronization issue

The pre-initialization CSA values also match the A2A3 sequential-reduction
reference exactly, connecting the model symptom to the selected ISA path.

The production fix still needs regression coverage for both platforms,
worker reuse and multiple kernel libraries, followed by model-suite tests.
This report validates the CPU simulator fix direction; A5 hardware bitwise
validation is outside its scope.