[aien-spark]

Draft. Review pending; do not add ready-to-merge.

Drake asked whether the CUDA++ attention library only takes the model shapes Aien runs. It did. This PR adds a general fallback kernel. The tuned kernels stay as the fast path, byte for byte.

Domain before and after (src/omega_gpu_attention_api.{c,h})

	before (main c30c0be)	after (this PR)
head_dim	64 or 128; anything else TOO_LARGE	1..256, any value (odd, 80/96/112, primes); > 256 TOO_LARGE
query heads	≤ 64	≤ 1024
GQA ratio	power of two only; 24q/1, 2, 4, 8 kv were refused with BAD_ARGS	any kv count dividing q heads (MHA, GQA of any ratio, MQA)
contiguous context	≤ 4096 (TOO_LARGE)	up to the 32-bit plane bound 2·kv·nkv·hd·elem ≤ 0xffffffff
mask	gqa: none; rows: causal	new omega_gpu_attention_general: causal or non-causal, left/right sliding window, q_pos0 offset, empty range writes zeros
scale	fixed 1/sqrt(hd)	0 means 1/sqrt(hd), otherwise the caller's value; negative or NaN is BAD_ARGS
KV	f32 contiguous / bf16 paged, dense layout	additionally f32 or bf16 with element strides (token, head)
softmax / accumulation	f32	f32 (online softmax over chunks of roundup(hd, 32))
RoPE	caller's	caller's: q and k are taken after rotary, as in FlashAttention

The existing entries (gqa_attention_f32, gqa_attention_rows_f32, paged_attention_bf16, _batch_bf16, reserve, reserve_rows) keep the tuned kernel for shapes it took before. Shapes it refused now run on the general kernel. The new entry omega_gpu_attention_general (EXPERIMENTAL) exposes the full domain.

Unchanged and still model-shaped:

• the fused served kernels omega_gpu_decode_attn (hd 128, 32/4), decode_attn256 (the live 35B), prefill_attn and prefill_attn256. This PR does not touch them.
• the general kernel's own limits:
  • scalar f32, no tensor cores;
  • one CTA per (row or sequence, head), so the paged general route has no context split;
  • no ALiBi, softcap or dropout.

Evidence (labels: EXERCISED on the host, chip-EXERCISED on GB10; nothing here is qualified)

Receipts are in evidence/ATTN-GENERAL-20261027-4d51f91f/ (SUMMARY.md).

• General sweep on chip: 318 of 318 seeded configs pass (seed 0xa77e0601), 348 checks with 0 failed, against an f64 host oracle.
  • Coverage:
    • hd {1, 2, 3, 5, 7, 31, 33, 63, 64, 65, 80, 96, 97, 112, 127, 128, 129, 160, 191, 192, 255, 256};
    • GQA ratios 1-16 and 1024;
    • kv 0..5000;
    • 115 non-causal configs and 67 windowed;
    • all five entries.
  • Worst error was 0.0062 of the tolerance.
  • Tolerance: rel·|want| + 2e-5, with rel 2e-4 (kv ≤ 2048) or 4e-4 (above).
    • WolframLanguageEvaluator: N[4096*2^-24] -> 2.44e-4
    • WolframLanguageEvaluator: N[5000*2^-24] -> 2.98e-4
• Mutants: all five caught on chip:
  • GQA_MAP at config 4;
  • TAIL_CHUNK at 4;
  • TAIL_DIMS at 2;
  • CAUSAL at 0;
  • SCALE at 2.
• Routing proof: a mutant planted in the general kernel shows when the general route is forced and is absent on the default tuned route.
• Refusal tests (--host-only):
  • hd 0 or 257, q heads 1025, kv not dividing q, negative or NaN scale, NULL pointers, plane overflow, bad window.
• Batteries on chip: hd 64 219/0, hd 128 225/0.
• Head sweep on chip: 72/72 points. It now includes the 8 non-power-of-two GQA points that were refused before.
• Host: make test reports CUDA++ HOST TESTS PASS (46 targets). make test-gpu-attention now also runs --general --sim.

Served shapes: no slower

• Static check: all 252 tuned kernels are byte-identical to main (omega_gpu_attention_codegen_hd, hd 64/128 × f32/bf16 × log2bs 0..8 × log2gqa 0..6). The hd 64 digest pins in the test still pass.
• A/B timing: the AB-PLAN sha256 went on BOARD before each timed run. The workload was main c30c0be vs this branch, tuned gqa_f32 32q/4kv hd 64 ctx 2048, run beside the live 35B. Samples overlapping 18091 in_flight > 0 were discarded.
  • v1, 10+10 runs: median ratio B/A 1.061, Mann-Whitney p 0.853. That is outside the predeclared 1.05 band, so it is reported as measured.
  • v2, 40 pairs (38 A / 40 B clean): WolframLanguageEvaluator Median[b]/Median[a] -> 0.9990, MannWhitneyTest p 0.960. Within the band.

Why C, and the .om successor

The kernels are emitted through the existing Blackwell IR builder inside the existing C file src/omega_gpu_attention_api.c. Device kernels stay on the encoder until GpuStack PLAN Track C1 (Rules lines 43-73), and this PR adds no new C file.

• The .om successor is the C1 attention lowering. The general kernel's semantics (range rule, chunked online softmax, GQA map) are written in the header so that port can be checked against the same sweep and receipt schema OMEGA_GPU_ATTENTION_GENERAL_V1.
• The code is libm-free (a Newton sqrt on the host side), and the file stays in the freestanding closure (portability row 59).

API

Runtime API 0.5.0 → 0.6.0 (MINOR, per API.md semver item 4):

• new struct OmegaGpuAttnGeneralDesc (64 B);
• new fn omega_gpu_attention_general (EXPERIMENTAL);
• new consts GENERAL_MAX_HEAD_DIM, GENERAL_MAX_Q_HEADS and WINDOW_NONE;
• 12 internal test hooks.

make test-runtime-api passes: 0.6.0, 192 INTERNAL.

#139 (GEMM) also bumps to 0.6.0. Whichever merges second takes 0.7.0, and it also rebases the README §10 line.

Not run:

• NOT_RUN: timing A/B for the paged bf16 and hd 128 tuned entries. Their kernels are byte-identical, and only the host routing changed.
• NOT_RUN: a tensor-core general kernel.