KDA Chunk Prefill on Ascend A3

Project objective

Implement efficient KDA chunk-prefill forward computation on Huawei Ascend A3 using EasyASC, and deliver a reusable model for evaluating how multiple kernels should be partitioned, organized, and scheduled. Compare at least two substantially different partitions: predict their resource use and performance, implement the selected design, and use experiments to verify or revise the predictions.

The task covers Q/K normalization and the KDA core. QKV projections, short convolution, raw gate activation, output gating/normalization, and output projection are outside the task. Model weights are not required. Backward is an optional extension; a token-by-token decode loop is not the target implementation.

Function and tensor contract

def forward(q, k, v, g, beta, initial_state=None, *, aux=None):
    return o, final_state

All dimensions are positive and HV%H==0. Value head h uses Q/K head h//(HV//H). Core cases use H=HV and K=V=128. The checker passes unmodified, unnormalized CPU tensors; the adapter orchestrates OpExec kernels and same-dtype CPU/NPU transfers. Return contiguous tensors and preserve all input values, dtypes, shapes, and auxiliary dictionary structure.

Mathematical definition

Normalize Q and K along their last dimension with the exact rounding order:

xf = float32(x)
xnorm = bfloat16(xf / sqrt(sum(xf*xf, dim=-1, keepdim=True) + 1e-6))

The reference converts the normalized bf16 values to fp32 for chunk computation. Epsilon is inside the square root; query scaling is fixed at K**-0.5. For one batch and value head, with state S shaped [K,V], the equivalent recurrent semantics are:

Sdecay = exp(g_t)[:,None] * S_previous
r_t = v_t - k_t^T @ Sdecay
S_t = Sdecay + beta_t * k_t[:,None] * r_t[None,:]
o_t = (K**-0.5 * q_t)^T @ S_t

Use chunk-local parallel computation and inter-chunk state carry in the submitted implementation. An equivalent mathematical construction for a chunk of C tokens is:

G_i = sum_{u=0..i} g_u
L_ij = beta_i * sum_d(k_i,d*k_j,d*exp(G_i,d-G_j,d)) if j<i else 0
A = solve(I+L, diag(beta))
W = A @ (exp(G) * k)
U = A @ v
Qg_i = K**-0.5 * q_i * exp(G_i)
Kg_i = k_i * exp(G_last-G_i)
Aqk_ij = K**-0.5 * sum_d(q_i,d*k_j,d*exp(G_i,d-G_j,d)) if j<=i else 0
Vnew = U - W @ S_in
O = Qg @ S_in + Aqk @ Vnew
S_out = exp(G_last)[:,None] * S_in + Kg.T @ Vnew

Indices are zero-based and ordered by token position. The golden uses chunk size 64, zero-padding q/k/v/g/beta inside the reference for the last partial chunk, trimming output while preserving state. Your entry receives the original T. Your internal chunk size, intermediate precision/layout, kernel count, and fusion boundaries are unrestricted, provided the contract and accuracy tests pass. The bundled golden.py and chunk implementation define numerical comparison behavior.

Example

Use a legal core shape: B=H=HV=1,T=128,K=V=128. At every token, q and k equal e0=[1,0,...,0], v is all ones, g is zero, beta is one, and initial_state is None. After normalization and bf16 rounding, q/k remain e0.

o: every element is 0.08837890625 = bfloat16(1/sqrt(128))
final_state: row 0 is all ones; all other rows are zero

The first token initializes row 0 of the state; subsequent tokens preserve it. Setting beta to zero with a zero initial state instead produces zero output and zero final state. These examples explain semantics; the coverage requirements below define the evaluation domain.

Auxiliary matrices and host policy

aux[L] supplies read-only contiguous bf16 matrices for L in {32,64,128}: lower[i,j]=1[j<=i], upper[i,j]=1[i<=j], ones[i,j]=1, and identity[i,j]=1[i==j]. Both triangular matrices include their diagonals. For [L,K] data, lower@X gives inclusive token-prefix sums. Such arithmetic belongs in submitted kernels.

Auxiliary setup is excluded from timing, but device loads, conversions, and use count toward kernel duration. The matrices depend only on dimensions, not on input values or golden results.

Allowed host actions are shape/dtype/device queries, metadata-only size arithmetic, element-order-preserving no-copy views/reshapes/flatten/squeeze/unsqueeze/detach, uninitialized empty-family allocation, and same-dtype CPU/NPU transfer. Numeric tensor computation, initialized tensor creation, dtype conversion, slicing/indexing/transposition/permutation, copying reshape, data extraction (item, tolist, numpy, storage/pointers), and reference/alternate-engine imports are forbidden. Normalization, cumsum, matmul, conversion, and layout work belong in kernels.

The guard audits submission import and execution, including local dependencies and participant callbacks. Each case must observe trusted OpExec runtime operations. Exceptions caught by a submission do not erase guard violations. There is no submission flag to disable the guard. This is a basic in-process audit, not proof of physical A3 execution or a hostile-code sandbox.

Public examples and full evaluation

For both outputs, simultaneously require abs(actual-golden)<=0.02+0.02*abs(golden) and relative L2 error at most 0.02 when the reference is nonzero. An all-zero reference uses absolute error only. Reject NaN/Inf, incorrect shape/dtype, non-contiguous output, missing final state, input mutation, or host violations. All instructor core cases must pass. The public checker is only a development check. Internal compute/storage precision is a design choice. A3-specific hardware calibration of the tolerance remains pending.

The table below lists only downloadable samples. The full evaluator retains the complete regression suite and additional seeds and tail lengths within the declared domain. The same golden, tensor checks, host policy and tolerances apply. Exact private instances are not distributed; tuples below are (B,T,H,HV,K,V).

The input generator uses a CPU torch.Generator seeded per row. Each raw Q/K uses gains exp(randn(B,T,H,1)*0.6), multiplied by an independently generated fp32 normal tensor, then converted to bf16. Tiny mode scales these values by 1e-5 and rounds back to bf16; zero_rows zeroes q at ::3 and k at 1::3. V is standard-normal rounded to bf16.

Gate values are logsigmoid(randn)/divisor, with divisors 128, 8, and 1 for weak/moderate/strong; zero mode uses zeros. Beta starts as sigmoid(randn) and is optionally set to zero, one, or the repeating edge values {0,1,1e-4,1-1e-4} along token positions. Nonzero state is fp32 standard-normal times 0.25; zero state is an explicit tensor; none passes None. Preserve generation order and the PyTorch version to reproduce values.

Environment and commands

Extract the student package and enter kda_a3_practice/. Reference tests require PyTorch, einops, and pytest; no FLA installation, model weights, Triton, CUDA, CANN, or NPU is required for CPU reference checks.

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt
python preflight.py --target cpu
python kda_a3_exercise/check.py --list
python kda_a3_exercise/check.py --check-golden --suite public
cp kda_a3_exercise/submission.py my_solution.py

CPU checks have a pinned environment in ENVIRONMENT.md. EasyASC simulator and physical A3 execution are unverified on this machine because the runtime and NPU are absent; python preflight.py --target easyasc or --target ascend records this explicitly. For implementation, install the EasyASC community package and its runtime dependencies. Set EASYASC_ROOT to the directory containing the easyasc/ Python package. Use from easyasc.a3 import *; one Python process should import only one device target.

export EASYASC_ROOT="$HOME/projects/easyasc"
export PYTHONPATH="$EASYASC_ROOT${PYTHONPATH:+:$PYTHONPATH}"
# Implement my_solution.py before running submission checks.
python kda_a3_exercise/check.py --submission my_solution.py --suite public
python kda_a3_exercise/check.py --submission my_solution.py --suite public --report public-correctness.json

Use --case public_513 to isolate a failure. Exit codes are 0 for all selected cases passing, 1 for case failures, and 2 for argument/loading errors. The unimplemented template fails explicitly. --check-golden verifies the reference harness, not an A3 submission; independent mathematical examples live in tests/. Reports always mark evaluation_scope=public_samples and full_evaluation=false.

Develop with OpExec(..., simulator=True) and enable a trace with trace="results/trace.json". A3/C220 exchanges between Cube and Vector through GM workspace; A5's L0C→UB and UB→L1 paths cannot be copied directly. Recheck L1/L0/UB resource limits and Pipe.S waiting behavior. Run A5 reference experiments in a separate process.

Physical execution uses simulator=False and a compatible CANN/SoC environment. Optional small-case device-code simulation uses CAModel, not cannsim record. Generate with OpExec(..., simulator=False, debug=True, gen_only=True), change -r npu to -r sim in both generated b.sh and r.sh, and select the actual A3 SoC, for example Ascend910_9362. CAModel is for small correctness checks, not performance reporting.

Timing and required submission

Compare the sum of actual device kernel durations, including normalization, conversions, and recomputation. Exclude host scheduling, launches, gaps between calls, and allocation. Preserve actual dependencies, execution order, and L2 reuse conditions. Report warmups, repetitions, cache conditions, units, and per-invocation timing.

Separate EasyASC simulator cycles from A3 device microseconds. The A3 simulator currently uses a2_cycle_model.json; its cycle estimates are not A3 hardware timings. The checker reports kernel_timing=null and device_execution_verified=false. Device execution and chunk-local parallelism require code review plus trace/device evidence.

Submit the A3 implementation and a reusable partition-evaluation model as code, a tool, a skill, or a document. For two different candidate partitions, include DAGs, per-kernel work and launch counts, chunk-local parallelism, state dependencies, L1/L0/UB peaks, buffer slots, workspace, L2 working set, logical GM traffic versus actual HBM traffic, intermediate dtypes, error budget, synchronization, predicted cost, and measured cost. Explain L2 residency rather than assuming every GM exchange reaches HBM. If a candidate is not implemented, justify rejecting it with resource calculations or a focused experiment.

Include the contract/case hashes, environment and SoC, host-audit report, passed/failed suites, correctness for every timed shape, traces, raw results, reproducible commands, and conclusions that generalize to other shapes. Use the included EVALUATION.md template. Backward remains optional: its reference loss is sum(o.float()*do.float())+sum(final_state*dht), returning gradients for q/k/v/g/beta/initial_state in their corresponding dtypes, with None for an absent initial state.