feat(zkLLM): Add proof-bound TensorBus execution for Llama-7B/2048

[hero78119]

Problem

Ceno needs a proof-bound tensor execution path for the zkLLM-compatible Llama-2-7B batch-one prefill workload at sequence length 2048. The existing scalar tensor path represents dense matrix work as one call per scalar output and materializes product-oriented witness data. That representation is not suitable for production-width Llama layers: it creates excessive ECALL volume, circuit width, and witness-memory pressure, while the GPU provider is forced into many small launches.

This PR adds the TensorBus resident execution model, row-batched MatMul proving, and the Llama-shaped attention/FFN interfaces needed to build the production workload. The production-layer base proof and independent Rust-verifier E2E remain the next validation gate.

Design Rationale

TensorBus and guest execution

Tensor intermediates are represented as opaque TensorBus handles rather than ordinary RISC-V heap or stack objects. Guest RAM is used only at explicit segment boundaries. A Llama-shaped guest can therefore execute a complete resident block as follows:

for layer in 0..32 {
  for head in 0..32 {
      // RAM → TensorState: WRITE hidden record
      let hidden = import(layer, Projection, &hidden_ram);

      // Internal TensorState:
      // READ hidden records; WRITE Q/K/V records
      let qkv = stage(layer, head, Projection, hidden);

      // READ Q/K/V records; WRITE score/probability/context records
      let context = stage(layer, head, Attention, qkv);

      // TensorState → RAM: READ context records
      export(layer, Attention, context, &mut context_ram[head]);
  }

  // RAM → TensorState: WRITE hidden + complete-context records
  let post_input = import(layer, PostFfn, &[&hidden_ram, &context_ram]);

  // READ hidden/context; WRITE next-hidden records
  let next_hidden = stage(layer, 0, PostFfn, post_input);

  // TensorState → RAM: READ next-hidden records
  export(layer, PostFfn, next_hidden, &mut hidden_ram);
}

The production topology uses IMPORT_BEGIN -> N x (attention, FFN) -> EXPORT_END for an atomic segment. The segment owns TensorRef versions, ordering, and residency. Attention outputs, residuals, norms, and FFN intermediates remain Tensor-space values; they do not become RISC-V memory. The planned production E2E also exercises smaller cache=1 attention pieces, where an attention group may be imported/exported between segments when required by the device budget.

Preservation of the Ceno proof protocol

The implementation is integrated into the existing Ceno Core, tower, batched-main sumcheck, transcript, witness-commitment, verifier, and recursion structure. MatMul contributes additional pointwise correction expressions to the unchanged batched-main reduction, following the existing prove_rotation-style claim flow. It does not add a tower layer, replace the main sumcheck, add PCS rounds, or introduce a second commitment system.

TensorBus boundary records remain explicit and independently constrained. Each TensorRef has one producer and one consumer edge, and each HintRef identifies a logical weight tile by (profile, layer, role, tile_index). Model weights are generated lazily and reused by logical identity; the guest has no production hint_base pointer.

Row-batched MatMul

A matrix operation T[M,K] x S[K,N] is represented by M compact row records. Records are keyed by (shard_id, segment_id, operation_ordinal, row_index), ordered without duplicates, and must cover every row 0..M. All rows share the operation, shape, TensorRef, and HintRef identities and write disjoint ranges of one output TensorRef. They are coalesced into one tiled GPU GEMM rather than launching one kernel per ECALL.

The reduction retains the exact signed-byte decomposition, range checks, canonical Q16/Q20 quotient/remainder relation, and product consistency. The auxiliary matrix sumcheck proves:

Q(output_point) * 2^shift + R(output_point)
  = sum_(k, section) A(k, section) * W(k, section)
      * eq(output_section_point, section)

Its claimed sum is checked against the expected sigma and produces terminal evaluations for A, W, Q, and R. Those evaluations and their points are appended to MatMul's first GKR layer and then fused into the existing batched-main expression. The original single WitIn PCS opening authenticates the resulting global evaluations.

This changes representation rather than arithmetic. One production layer still performs 448,824,082,432 matrix multiplications, but the projected scalar-output call count falls from 229,638,144 to 145,408 row records. The 32-layer pass plus LM head falls from 7,413,956,608 scalar calls to 4,655,104 row calls. K and N therefore increase MLE height and matrix-sumcheck work instead of multiplying physical circuit columns and materialized product witnesses.

Llama-2-7B/2048 workload

The target is batch-one full-sequence prefill, not one-token decode. Each layer uses hidden=4096, heads=32, head_dim=128, intermediate=11008, and exact causal attention over the full 2048 x 2048 query/key domain. The workload includes RMSNorm, Q/K/V/O projections, RoPE, causal masking, five lookup-based softmax relations, P x V, residuals, gate/up/down projections, SwiGLU, and Q16/Q20 rescaling. The complete pass additionally includes embedding, final RMSNorm, the 4096 x 32000 LM head, and lowest-index argmax.

Private weights are deterministic unauthenticated witnesses. Computation and workload are proof-bound to the shard witness commitment, but the proof does not authenticate a public Llama model root.

Residency and performance trade-off

Cache level 1 is the default and keeps reusable Tensor/RMM data resident for the guest execution path. Cache level 0 remains the existing on-demand host-resident path. No PIOP cache or new general-purpose memory manager is introduced.

Production attention is evaluated in independently admitted segments so that segment-owned QK, softmax, PV, and temporary witness MLEs can be released before the next segment. The 4090 target uses a separate cache=1 guest configuration with a 4--6 GiB base-resident set and a four-head fused QK/softmax/PV piece targeting approximately 8--10 GiB of active traces. The 5070 Ti target uses the same cache=1 model with one- or two-head pieces when its smaller budget requires them. Additional explicit guest-RAM import/export is accepted for these smaller pieces; the proof protocol is unchanged.

Change Highlights

• ceno_emul: TensorBus handles, TensorRef/HintRef tracking, resident providers, Llama-shaped CPU/GPU reference paths, and production tensor syscall plumbing.
• ceno_zkvm: TensorBus and production Core registration, row-batched MatMul, attention/softmax/boundary chips, GPU assignment, and native verification integration.
• ceno_recursion_v2: replay of the MatMul correction expressions through the existing batched-main and recursion path.
• ceno_rt and examples: guest-facing tensor-handle APIs plus tiny, resident-block, topology, MatMul, and production-layer guest programs.
• tensor-vm.md: workload definition, TensorBus semantics, MatMul reduction, residency model, and proof-boundary rationale.

Testing

The E2E plan is staged so that base proving and independent Rust verification precede recursion:

1. llama-tiny chained two-layer and four-layer resident guests: GPU output/reference equality, real base proof, native Rust verification, and rejection of TensorRef, HintRef, matrix-product, quotient, remainder, and main-claim tampering.
2. Production attention-stage guests: cache=1 one-/two-head and four-head segment variants, exact causal attention and lookup relations, legal import/export boundaries, and measured peak VRAM/host memory and H2D/D2H traffic.
3. One complete production Llama-2-7B/2048 layer: Q/K/V/O, RoPE, causal QK, softmax, P x V, both RMSNorms, residuals, SwiGLU, lazy HintRefs, real GPU base proof, and independent Rust verification.
4. Full 32-layer prefill plus LM head: legal segment packing, minimum successful shard count, final argmax, and recursion replay after the base proof gate.

Validated in the current implementation:

• Tiny two-layer and four-layer resident GPU guests produce independently verified base proofs and reject the listed tamper cases.
• Production-width provider preflight completes with 1,552 dense K1024-by-N128 panels and the static guest path avoids the former full-tensor initialization loop.
• Production-layer base proof and independent Rust verification are not claimed complete by this PR yet.

Risks and Rollout

• Production-width witness residency can exceed local GPU capacity. Segment admission must fail with an allocation ledger rather than silently changing the workload or proof protocol.
• Cache=1 retains reusable data by design; active segment bytes and base-resident bytes must be reported separately to distinguish intentional residency from leaks.
• The private-weight path proves computation relative to committed witnesses but does not provide public model authentication.
• Existing pointer-ABI tensor calls and non-tensor RISC-V memory behavior remain unchanged. The handle ABI is opt-in for generated tensor guests.

Follow-ups

• Complete the production-layer base-proof and independent-verifier E2E, then run recursion-v2 replay.
• Complete the full 32-layer segment/shard packing study on the 5070 Ti and RTX 4090.
• Add benchmark tables and throughput comparisons after the correctness and residency gates are complete.
• Public model-weight authentication, one-token decode, KV-cache decode hints, and authenticated persistent KV cache remain outside this PR.

Copilot Reviewer Directive (keep this section)

When Copilot reviews this PR, apply .github/copilot-instructions.md strictly.

[hero78119]

Fused attention proof timing and size projection

Assumptions:

• One attention layer contains 32 attention-head shards and one post-FFN shard.
• A 32-layer model contains 1,024 attention-head shards and 32 post-FFN shards (1,056 shards total).
• Single-head fused timings are measured on an RTX 5070 Ti. Layer, model, and multi-GPU values are projections unless stated otherwise.

create_proof timing

Phase	Head x1 (measured)	Attention x1: 32 heads (projected)	32 layers (projected)
Commit	0.444 s	14.208 s	454.656 s
Chip/tower scheduling	0.615 s	19.680 s	629.760 s
Main sumcheck	0.619 s	19.808 s	633.856 s
PCS opening	0.421 s	13.472 s	431.104 s
Attention create_proof	2.099 s	67.168 s	2,149.376 s
Post-FFN create_proof	-	1.369 s	43.808 s
Total create_proof	2.099 s	68.537 s	2,193.184 s (36.553 min)

The fused attention and PV tasks overlap inside the chip/tower scheduling interval and are not added separately.

Witness generation is excluded from all proving-time estimates. The current AOT/preflight witness path is not treated as the long-term execution architecture; witness generation is expected to use the multi-GPU PyTorch path independently of proving.

Serialized proof size

Scope	Size
Head x1	1,481,702 B (1.413 MiB)
Attention x1: 32 heads	47,414,464 B (45.218 MiB)
32-layer attention: 1,024 heads	1,517,262,848 B (1.413 GiB)

The 32-layer size excludes the 32 post-FFN proofs because their individual serialized size was not recorded. For reference, the measured pre-fusion 33-shard layer proof was 49,070,606 B; multiplying it by 32 gives 1,570,259,392 B (1.462 GiB), including post-FFN proofs but using the previous attention architecture.

16 x RTX 5070 Ti projection

An even split assigns 66 shards to each card: 64 attention-head shards and two post-FFN shards.

Per-card critical-path workload	Time
Commit	28.416 s
Chip/tower scheduling	39.360 s
Main sumcheck	39.616 s
PCS opening	26.944 s
Attention proof subtotal	134.336 s
Post-FFN proof: two shards	2.738 s
create_proof subtotal	137.074 s
Ideal 16-card proving time	137.074 s (2 min 17 s)

The ideal figure includes only create_proof work and assumes perfect shard parallelism with no shared CPU, storage, PCIe, or orchestration bottleneck. Witness generation, AOT/preflight, preparation, key generation, and independent verification are excluded.