[WS1 Test]  #154 dtype pin + #152 KV path + #153 bwd invariance + #150 chain gate

## Summary

Once WS1 operator implementations and ground-truth paths (#108, #145–#151) are largely in place, remaining work is no longer “write another op.” It is to close **contract axes, path consistency, cross-op backward audit, and the end-to-end exit gate** before WS2.

This issue is the **unified description and acceptance entry point** for that closing work. Execution is owned by four existing issues:

| Issue | Title | Role | Close order |
|-------|-------|------|-------------|
| [#154](https://github.com/RL-Align/RL-Kernel/issues/154) | dtype coverage in the numerical contract | Pin the dtype axis of the #108 contract | **First** (blocking pin) |
| [#152](https://github.com/RL-Align/RL-Kernel/issues/152) | KV-cache path consistency (prefill & decode) | Rollout (decode) vs training (prefill) numerical parity | After #154; **parallel** with #153 |
| [#153](https://github.com/RL-Align/RL-Kernel/issues/153) | Backward-pass consistency across all ops | Cross-op gradient batch invariance | After #154; **parallel** with #152 |
| [#150](https://github.com/RL-Align/RL-Kernel/issues/150) | Chain integration test (WS1 exit gate) | End-to-end forward + backward exit gate | **Last** |

---

## Background and motivation

1. **“Aligned” is meaningless without a pinned dtype set.** RL training runs in BF16; FP32 is the tolerance / gold reference. The dtype axis must be locked first, or each op validates under a different set and chain-level guarantees are empty.
2. **Rollout uses decode; training uses prefill.** The same token can get different logprobs under different reduction orders. Decode must be covered explicitly—not only via chunked-prefill.
3. **Forward alignment is not enough.** Gradient reductions that drift with batch shape break training. Backward must be a first-class tested property.
4. **All ops green ≠ full chain invariant.** Drift appears at op boundaries, dtype handoffs, and chunked-prefill / padding interactions. #150 is the real WS1 exit gate.

### Sole definition of “aligned”

All pass/fail decisions use **only** the #108 harness and `tolerance_contract.json` (extended by #154). No private threshold tables.

The contract must distinguish four judgments:

| Judgment | Meaning |
|----------|---------|
| `forward_accuracy` | Candidate vs strict FP32 / reference |
| `forward_invariance` | Same-dtype output invariance across batch / layout / path configs |
| `gradient_accuracy` | Candidate VJP/grad vs reference |
| `gradient_invariance` | Same-dtype gradient invariance across equivalent configs |

Gradient thresholds **must not** silently inherit forward thresholds.

---

## Goals

- [ ] Finalize the **dtype axis** of the numerical contract so every WS1 op validates against the same pinned set.
- [ ] Prove **prefill ↔ decode** (including generate-then-re-score) numerical consistency on the KV-cache path.
- [ ] Make **batch-invariant backward** a reusable harness property required for every WS1 op.
- [ ] Assemble the full standard-Transformer chain and pass the **WS1 exit gate** under the #108 contract.

## Non-goals (explicitly out of scope)

- Multi-GPU / TP / SP / NCCL (WS2)
- Real vLLM vs real Megatron alignment (WS3)
- Layer-wise probe / drift dashboard as a product (WS4)
- FP8, MoE, non-standard attention, throughput benchmarking
- Re-implementing individual op kernels (owned by #145–#151)

---

## Ownership boundary

| Role | Owns | Does not own |
|------|------|--------------|
| #108 + #145–#151 owners | Merged, callable reference / candidate paths | Contract / gate implementation details in this package |
| #154 | Dtype axis, four-judgment schema, one-time accumulation/TF32 pin | Inventing per-op thresholds |
| #152 | Prefill/decode path consistency, B1 stateful cache path, generate-re-score | Attention internal accumulation design (#147) |
| #153 | Reusable gradient harness, audit matrix, defect routing | Fixing red bwd kernels (push back to owning op issues) |
| #150 | Full-chain assembly, config sweep, first-drift diagnostics, CI gate | Force-closing while the handoff matrix is still red |

**Per-row readiness for the #150 gate:**

Merged implementation · registered reference/candidate · forward compliant with #154 · backward compliant with #153 where applicable · unresolved failures have a linked blocker

---

## Dependencies and execution order

```text
PHASE 2a (blocking)     #154 dtype / four-judgment schema
              │
              ├────────────────────┐
              ▼                    ▼
PHASE 2b      #152 KV path    #153 bwd harness
              │                    │
              └─────────┬──────────┘
                        ▼
PHASE 4       #150 full-chain integration = WS1 EXIT → WS2
```

**Close dependency:**

```text
             ┌→ #152 ─┐
#154 merged ─┤        ├→ #150
             └→ #153 ─┘
```

| Constraint | Meaning |
|------------|---------|
| Single source of truth | #108 harness + tolerance contract only |
| Ops are inputs | #145–#151 own kernels; #150–#154 own integration/contract axes |
| #150 is last | Do not close the exit gate while #152/#153 (or dtype) are red |

Live status lives only in [`ws1-handoff-matrix.md`](./ws1-handoff-matrix.md). When a row flips green/red, leave a short comment on #202 and #150 with an evidence link.

---

## Workstream 1 — #154 dtype coverage (blocking pin)

**Assignees:** @Flink-ddd @maxiaosong1124 @EthanZero2Hero  
**Suggested duration:** 2–4 days · 3–5 small PRs

### Scope

| Axis | Decision |
|------|----------|
| Mandatory | **BF16** (RL training dtype) |
| Reference | **FP32** (tolerance / gold) |
| Optional | **FP16** (test if backend supports; otherwise explicit skip) |
| Out | **FP8** (out of scope this month) |
| Accumulation | BF16 inputs → **FP32 accum** (one rule for all ops) |
| TF32 | **Pinned once** in the contract; disabled on the reference path |

### Planned PRs

| ID | Description | Exit signal |
|----|-------------|-------------|
| 154-PR1 | Contract docs: dtype-axis prose | Names BF16 mandatory / FP32 reference / FP8 out |
| 154-PR2 | `tolerance_contract.json` four judgments + per-dtype rows | Table is the single source of truth |
| 154-PR3 | FP32-for-BF16 accum + TF32 decision | Op issues only reference, never re-define |
| 154-PR4 | Test helpers / `check_operator.py` use the same dtype set | Every WS1 op sweep is consistent |
| 154-PR5 | CI enforces mandatory BF16; optional skips are explicit | No silent fallback counted as pass |

### Acceptance

- [ ] Contract names the dtype set (BF16 mandatory, FP32 reference, FP8 out)
- [ ] All four judgments have per-dtype threshold rows as SSOT
- [ ] Every WS1 op uses the same pinned set
- [ ] Accum + TF32 rules are written once and referenced by op issues
- [ ] A backend claiming WS1 compliance must run mandatory BF16; skip/fallback must not be reported as a BF16 pass

### Repo docking

- `rl_engine/kernels/gtest/tolerance_contract.json`
- `rl_engine/kernels/gtest/tolerance.py`
- `tests/test_tolerance_contract.py`

---

## Workstream 2 — #152 KV-cache path consistency

**Assignee:** @zhangj1an  
**Related PR:** [#178](https://github.com/RL-Align/RL-Kernel/pull/178) (WIP)  
**Suggested duration:** ~1 week (including #178 close-out + generate-rescore + CI smoke)

### Scope

Ensure prefill and decode produce the same reductions for the same effective context:

1. Decode (one query vs N cached KV) uses the same fixed reduction order as the equivalent full prefill
2. Decode appears **explicitly** in the test sweep—not only chunked-prefill
3. Cache write/read layout and stored-KV dtype do not introduce writer/reader precision differences
4. **Generate-then-re-score** matches final selected logprobs (same token / position / mask / weight identity)

### Stateful validation levels

| Level | Path | Requirement |
|-------|------|-------------|
| A | Full prefill vs concat-based decode reference | Mathematical parity |
| **B1 (required to close #152)** | In-repo allocate → write → read → decode | Prove storage dtype / layout / mutation |
| B2 | Bind registered production KV candidate | Required when available; if absent, record explicitly—no production-alignment claim |

> `NativeKVCacheAttnOp` is a stateless mathematical reference and **does not by itself satisfy B1**.

### Planned PRs

| ID | Description |
|----|-------------|
| 152-PR1 | Full-prefill vs chunked-prefill consistency |
| 152-PR2 | Decode-stage path test (1 query vs N KV) |
| 152-PR3 | B1 stateful writer/reader; bind B2 when available |
| 152-PR4 | Generate-then-re-score selected-logprob equivalence |
| 152-PR5 | CI-friendly decode smoke (short / long / varlen / padded) |
| 152-PR6 | Remove private KV tolerances; use shared contract resolver only |

### Acceptance

- [ ] For a fixed sequence, decode logprobs match prefill logprobs within #108 tolerance (ideally bitwise)
- [ ] Decode is exercised directly in the sweep (batch=1/N, multiple padding layouts)
- [ ] Stored-KV dtype/layout adds no extra drift
- [ ] Level A + mandatory B1 pass; B2 passes when a production candidate exists, otherwise absence/blocker is explicit with no production-alignment claim
- [ ] Generate-then-re-score compares **final selected logprobs** (same token IDs, positions, masks, pre-update weights)
- [ ] short / long / varlen / padded covered; tolerances come only from the shared contract

### Repo docking

- `rl_engine/kernels/ops/pytorch/attention/kv_cache.py` (Level A only)
- `tests/test_kv_cache_attention.py`
- Add `rl_engine/testing/kv_consistency.py` (B1 entry point)
- PR #178

---

## Workstream 3 — #153 Backward-pass consistency

**Assignees:** @maxiaosong1124 @EthanZero2Hero @Flink-ddd  
**Suggested duration:** 1–1.5 weeks (harness / audit / matrix / CI; kernel fixes tracked under #145–#151)

### Scope

Make batch-invariant backward a required, tested property of every WS1 op:

1. Provide a reusable **gradient-invariance** assertion in the #108 harness (independent `gradient_*` rows)
2. Each op’s grad reduction (`dx` / `dweight` / `dW` …) uses a **fixed, batch-shape-independent** order — **no `atomicAdd` in backward**
3. Prioritize reduction-heavy paths: RMSNorm `dweight` → GEMM `dW`/`dX` → attention bwd → embedding-grad scatter → remaining ops
4. Sweep: batch=1/N, chunked-prefill on/off, padding layouts

### Logical-workload identity (must hold before comparison)

- Preserve sample / token / position / mask / weight logical identity and order
- Exclude padding tokens from output / loss / grad comparisons
- Restore packed/chunked results to canonical `(sample_id, token_position)` before compare
- Use identical upstream grads, loss reduction, and active-token denominator
- Compare parameter gradients over the **same complete logical sample set**, not 1-sample vs N-sample aggregate

### Planned PRs

| ID | Description |
|----|-------------|
| 153-PR1 | Logical-workload canonicalization + reusable gradient-invariance assertion |
| 153-PR2 | Register all WS1 refs/candidates and differentiable inputs/params |
| 153-PR3 | Gradient-diff reporting (max abs/rel, tensor name, first failing op) |
| 153-PR4 | no-`atomicAdd` / fixed-order rule (docs + checklist + optional guards) |
| 153-PR5 | Per-op forward+backward status matrix; wire full-chain bwd into CI (feeds #150) |

**Target harness API:**

```text
assert_gradient_batch_invariant(op, configs, contract) -> GradientInvarianceReport
```

### Acceptance

- [ ] #108 provides a reusable gradient-invariance assertion any op issue can call
- [ ] All WS1 refs/candidates, differentiable inputs/params, and config transforms are registered; accuracy and invariance are reported separately
- [ ] Every WS1 op with backward passes the full sweep (bitwise or within #108 tolerance)
- [ ] No op uses `atomicAdd` or shape-dependent accumulation in backward
- [ ] Short matrix records per-op fwd/bwd status; drift reports include max abs/rel, tensor name, first failing op

### Repo docking

- `rl_engine/kernels/gtest/op_checks.py` (`check_grad` → cross batch-config invariance)
- `operator_specs.py` `grad_input_names`
- `tolerance_contract.json` `gradient_*` rows

---

## Workstream 4 — #150 Chain integration test (WS1 exit gate)

**Assignees:** @maxiaosong1124 @EthanZero2Hero @Flink-ddd  
**Suggested duration:** 1–1.5 weeks (assembly + sweep + diagnostic + CI; buffer for cross-op boundary drift)  
**Close condition:** All hard deps green with handoff-matrix evidence — **last WS1 issue to close**

### Scope

Assemble all WS1 ops into one end-to-end forward + backward pass and assert invariance across the config matrix.

**Standard Transformer forward topology** (CI tiny model may shrink layers/hidden size, but topology and semantics must not be cut):

```text
token embedding
  → RMSNorm → Q/K/V → QK-Norm (if required by fixed model) → RoPE
  → attention → out proj → residual
  → RMSNorm → MLP gate/up → SwiGLU → down → residual
  → final RMSNorm → LM head → selected logprob → masked loss
```

**Config sweep:**

| Axis | Comparison |
|------|------------|
| batch | 1 vs N |
| chunked-prefill | on vs off |
| padding | multiple layouts |
| packing | packed vs unpacked where supported |
| dtype | #154 pinned set (BF16 e2e minimum; FP32 reference path) |

### Planned PRs

| ID | Description |
|----|-------------|
| 150-PR1 | Pin fixed workload manifest + topology-correct BI forward chain |
| 150-PR2 | Full backward through the same chain |
| 150-PR3 | batch / chunked-prefill / padding / packing sweep |
| 150-PR4 | First-drift diagnostic (max abs/rel + first failing op) |
| 150-PR5 | CI small model + GPU/backend provenance as WS1 regression gate |

**Suggested entry points:**

```bash
python scripts/ws1_chain_gate.py --model tiny-bi-transformer --dtype bfloat16 --seed 0
# or
pytest tests/test_ws1_chain_integration.py -q
```

Gate report must include: `requested_backend` / `actual_backend`, device/CC, dtype, model/workload ID, seed, config transform, fallback reason, max abs/rel, first failing layer/op.  
**Silent fallback to reference = fail** (unless the contract explicitly permits and labels it). A CPU-only green run **cannot** certify the GPU WS1 gate.

### Hard dependencies (all required before closing #150)

```text
#108 harness + contract (including #154 dtype axis)
#145 RMSNorm  #146 GEMM  #147 Attn  #148 Logprob
#149 RoPE/elementwise  #151 Embedding + LM head
#152 KV prefill/decode
#153 bwd harness + each op bwd green
        └──→ #150 EXIT
```

### Acceptance

- [ ] End-to-end logprob for a fixed prompt set is invariant across batch / chunked-prefill / padding (#108 tolerance)
- [ ] Full-chain gradients are invariant across the same matrix
- [ ] First failing layer/op can be identified (max abs / relative)
- [ ] Single-command CI-suitable gate; any op still failing in-chain **links back** to its owning issue
- [ ] Mandatory CUDA BF16 run records requested/actual backend; silent skip/fallback is not a pass
- [ ] Close only after every required handoff-matrix row is green with evidence

---

## Suggested schedule

| Window | Focus | Deliverable |
|--------|-------|-------------|
| **W0** | #154 schema + dtype pin | Four-judgment shape + mandatory BF16 policy merged |
| **W0–W2** | **Parallel** #152 (#178 + A/B1/B2 + generate-rescore) and #153 (harness + matrix + heavy-op audit) | Prefill/decode green; importable grad assertion; handoff matrix live |
| **W2–W3** | #150 assembly + sweep + first-drift + CI | Single-command e2e green = **WS1 exit** |

Parallelism notes:

- Pin #154 early; it barely blocks kernel work
- After #154 is stable, #152 ∥ #153
- #150 assembly may start once interfaces stabilize, but **must not close early**

---

## Cross-issue Definition of Done

### #154

- [ ] Contract states BF16 mandatory / FP32 reference / FP8 out
- [ ] Four-judgment per-dtype threshold table is SSOT
- [ ] Accum + TF32 pinned once
- [ ] All WS1 op tests use the same dtype set
- [ ] CI: mandatory BF16 always runs; optional capabilities skip explicitly; no silent fallback pass

### #152

- [ ] Prefill vs chunked-prefill consistent
- [ ] Decode explicitly tested and matches prefill within #108 tolerance
- [ ] Stored-KV layout/dtype adds no extra drift
- [ ] Mandatory B1 passes; B2 passes when available or is recorded absent without a production-alignment claim
- [ ] Generate-then-re-score selected-logprob equivalence holds
- [ ] Tolerances come only from the shared contract; requested/actual backend reported
- [ ] short / long / varlen / padded + CI smoke

### #153

- [ ] Reusable gradient-invariance assertion landed
- [ ] Logical-workload canonicalization enforced; refs/candidates and differentiable params registered
- [ ] Every WS1 op with backward passes the full sweep
- [ ] No `atomicAdd` / no shape-dependent bwd accumulation
- [ ] Fwd+bwd status matrix + standard drift report

### #150

- [ ] Full standard-Transformer fwd+bwd assembled on the BI path
- [ ] Workload manifest + topology pinned
- [ ] logprob + gradients invariant across batch / chunked-prefill / padding (/ packing)
- [ ] First-drift diagnostic works
- [ ] Single-command CI gate; failures link back to owning op issues
- [ ] CUDA BF16 records backend and rejects silent fallback
- [ ] All hard deps green + evidence → **last issue closed**

---

## Risks and external phrasing

| Risk | Handling |
|------|----------|
| #154 is light but blocking | Pin early or every later suite reworks dtype coverage |
| #152 is the most training-relevant WS1 check | Prioritize over more isolated unit tests |
| #153 does not replace op fixes | Harness green + matrix red → push back to #145/#146/#147/… |
| #150 is serial | Do not force-close while #152/#153 are red |
| What all four green actually prove | **Operator-level single-GPU train-inference self-consistency only** — not real vLLM == Megatron (that is WS3) |
---

Issue	Title	Role	Close order
#154	dtype coverage in the numerical contract	Pin the dtype axis of the #108 contract	First (blocking pin)
#152	KV-cache path consistency (prefill & decode)	Rollout (decode) vs training (prefill) numerical parity	After #154; parallel with #153
#153	Backward-pass consistency across all ops	Cross-op gradient batch invariance	After #154; parallel with #152
#150	Chain integration test (WS1 exit gate)	End-to-end forward + backward exit gate	Last

Judgment	Meaning
`forward_accuracy`	Candidate vs strict FP32 / reference
`forward_invariance`	Same-dtype output invariance across batch / layout / path configs
`gradient_accuracy`	Candidate VJP/grad vs reference
`gradient_invariance`	Same-dtype gradient invariance across equivalent configs

Role	Owns	Does not own
#108 + #145–#151 owners	Merged, callable reference / candidate paths	Contract / gate implementation details in this package
#154	Dtype axis, four-judgment schema, one-time accumulation/TF32 pin	Inventing per-op thresholds
#152	Prefill/decode path consistency, B1 stateful cache path, generate-re-score	Attention internal accumulation design (#147)
#153	Reusable gradient harness, audit matrix, defect routing	Fixing red bwd kernels (push back to owning op issues)
#150	Full-chain assembly, config sweep, first-drift diagnostics, CI gate	Force-closing while the handoff matrix is still red

Constraint	Meaning
Single source of truth	#108 harness + tolerance contract only
Ops are inputs	#145–#151 own kernels; #150–#154 own integration/contract axes
#150 is last	Do not close the exit gate while #152/#153 (or dtype) are red

Axis	Decision
Mandatory	BF16 (RL training dtype)
Reference	FP32 (tolerance / gold)
Optional	FP16 (test if backend supports; otherwise explicit skip)
Out	FP8 (out of scope this month)
Accumulation	BF16 inputs → FP32 accum (one rule for all ops)
TF32	Pinned once in the contract; disabled on the reference path

ID	Description	Exit signal
154-PR1	Contract docs: dtype-axis prose	Names BF16 mandatory / FP32 reference / FP8 out
154-PR2	`tolerance_contract.json` four judgments + per-dtype rows	Table is the single source of truth
154-PR3	FP32-for-BF16 accum + TF32 decision	Op issues only reference, never re-define
154-PR4	Test helpers / `check_operator.py` use the same dtype set	Every WS1 op sweep is consistent
154-PR5	CI enforces mandatory BF16; optional skips are explicit	No silent fallback counted as pass

Level	Path	Requirement
A	Full prefill vs concat-based decode reference	Mathematical parity
B1 (required to close #152)	In-repo allocate → write → read → decode	Prove storage dtype / layout / mutation
B2	Bind registered production KV candidate	Required when available; if absent, record explicitly—no production-alignment claim

ID	Description
152-PR1	Full-prefill vs chunked-prefill consistency
152-PR2	Decode-stage path test (1 query vs N KV)
152-PR3	B1 stateful writer/reader; bind B2 when available
152-PR4	Generate-then-re-score selected-logprob equivalence
152-PR5	CI-friendly decode smoke (short / long / varlen / padded)
152-PR6	Remove private KV tolerances; use shared contract resolver only

ID	Description
153-PR1	Logical-workload canonicalization + reusable gradient-invariance assertion
153-PR2	Register all WS1 refs/candidates and differentiable inputs/params
153-PR3	Gradient-diff reporting (max abs/rel, tensor name, first failing op)
153-PR4	no-`atomicAdd` / fixed-order rule (docs + checklist + optional guards)
153-PR5	Per-op forward+backward status matrix; wire full-chain bwd into CI (feeds #150)

Axis	Comparison
batch	1 vs N
chunked-prefill	on vs off
padding	multiple layouts
packing	packed vs unpacked where supported
dtype	#154 pinned set (BF16 e2e minimum; FP32 reference path)

ID	Description
150-PR1	Pin fixed workload manifest + topology-correct BI forward chain
150-PR2	Full backward through the same chain
150-PR3	batch / chunked-prefill / padding / packing sweep
150-PR4	First-drift diagnostic (max abs/rel + first failing op)
150-PR5	CI small model + GPU/backend provenance as WS1 regression gate

Window	Focus	Deliverable
W0	#154 schema + dtype pin	Four-judgment shape + mandatory BF16 policy merged
W0–W2	Parallel #152 (#178 + A/B1/B2 + generate-rescore) and #153 (harness + matrix + heavy-op audit)	Prefill/decode green; importable grad assertion; handoff matrix live
W2–W3	#150 assembly + sweep + first-drift + CI	Single-command e2e green = WS1 exit

Risk	Handling
#154 is light but blocking	Pin early or every later suite reworks dtype coverage
#152 is the most training-relevant WS1 check	Prioritize over more isolated unit tests
#153 does not replace op fixes	Harness green + matrix red → push back to #145/#146/#147/…
#150 is serial	Do not force-close while #152/#153 are red
What all four green actually prove	Operator-level single-GPU train-inference self-consistency only — not real vLLM == Megatron (that is WS3)

Uh oh!

[WS1 Test] #154 dtype pin + #152 KV path + #153 bwd invariance + #150 chain gate #244

Description

Summary

Background and motivation

Sole definition of “aligned”

Goals

Non-goals (explicitly out of scope)

Ownership boundary

Dependencies and execution order

Workstream 1 — #154 dtype coverage (blocking pin)

Scope

Planned PRs

Acceptance

Repo docking

Workstream 2 — #152 KV-cache path consistency

Scope

Stateful validation levels

Planned PRs

Acceptance

Repo docking

Workstream 3 — #153 Backward-pass consistency

Scope

Logical-workload identity (must hold before comparison)

Planned PRs

Acceptance

Repo docking

Workstream 4 — #150 Chain integration test (WS1 exit gate)

Scope

Planned PRs

Hard dependencies (all required before closing #150)

Acceptance

Suggested schedule

Cross-issue Definition of Done

#154

#152

#153

#150

Risks and external phrasing

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions