[WS2] Framework for ablation experiment

[CyberSecurityErial]

How to contribute

Below are areas that we have not yet created a dedicated knob to test whether the different implementation between training (megatron) and inference (vllm) caused the discrepancy in outputs.

If you want to claim for tasks in this PR, please comment below, as it is mainly reserved for internal collaborators.

You should base your PR on #230.

1. create a dedicated knob to turn on or off the listed feature.
2. add a custom written kernel that will be called when the knob is turned off.

GEMM knobs

Collected from @frank-2077's study of how vLLM and Megatron implements GEMM kernel differently.

#	Axis	What it isolates	Knob	Contract field	Binding tier	Runnable today?	Task	Owner
1	Runtime flags	Whether drift is arithmetic at all, or just two differently-configured processes	✅ training.deterministic_mode, rollout.batch_invariant, gemm.compute_dtype, gemm.tf32, gemm.bf16_reduced_precision_reduction, training.sequence_parallel, training.tensor_parallel_size	❌ no GemmDeterminismSpec	IDENTICAL (dtype, TP/SP) + SEMANTIC (TF32, reduced-precision reduction) + RECORDED (the rest)	⚠️ probe-only; nothing rejects a mismatch	extend determinism.py to the GEMM-relevant flags and fail closed on the blocking ones	🙋
2	GEMM kernel	Megatron cuBLAS vs. vLLM Triton vs. fixed-order det_gemm	✅ gemm.kernel	❌ no GemmContract.backend_id	RECORDED (measured, never forced equal)	❌ neither side is dispatchable through the catalog	GemmBackendCapability + the gemm.kernel knob; #239 PR2 the shard reference to compare against	🙋
3	ColumnParallel backward dX all-reduce	Whether the TP dX sum, not the GEMM, moves the number	✅ gemm.dgrad_reduce	❌ no GemmReductionSpec	SEMANTIC (order, acc_dtype) + RECORDED (engine)	❌ no fixed-order reference reduction exists	#239 PR5 the fixed-order reference; #239 PR6 the dX communication operator; the contract field	🙋
4	RowParallel forward all-reduce / reduce-scatter	Whether SP reduce_scatter and non-SP all_reduce are the same number	✅ gemm.forward_reduce	❌ no GemmReductionSpec	SEMANTIC (order, acc_dtype) + RECORDED (collective, engine)	❌ the reduce_scatter arm has no operator	#239 PR3 the forward AR operator; #247 the ReduceScatter kernel; the contract field	🙋
5	Fused wgrad accumulation	Whether fusing dW = dY^T X with += main_grad changes the result	✅ gemm.wgrad_accumulate	❌ no GemmWgradSpec	SEMANTIC (accum_dtype) + RECORDED (fusion)	❌ the unfused reference exists only as det_gemm_db, with no accumulation wrapper	#239 PR7 the dW operator + explicit-accumulation reference; the contract field	🙋
6	vLLM all-reduce backend	Whether vLLM's handwritten IPC / MNNVL kernel differs from NCCL	✅ rollout.all_reduce_backend	❌ no field; recorded via rollout_recorded_extra	RECORDED (caller-supplied extra)	⚠️ selectable in vLLM, not declarable in the contract	read the actually selected backend in VllmProvenanceAdapter; #239 PR3 compare against the NCCL and reference arms	🙋

Attention knobs

Collected from #235 Attention contract work and the Megatron/vLLM implementation study.

The user-facing knob should remain one top-level attention mismatch knob. The rows below are the dimensions that attention=on must record and isolate in the mismatch report; they are not necessarily separate user-facing knobs. When an optimized path is disabled for attribution, the fallback must be an RL-Kernel deterministic/custom reference path, not generic PyTorch SDPA.

#	Axis	What it isolates	Knob	Contract field	Binding tier	Runnable today?	Task	Owner
1	Attention runtime/provenance	Whether drift is real arithmetic drift, or just two attention processes configured differently	✅ top-level attention; recorded fields attention.requested_backend, attention.actual_backend, attention.mode, attention.fallback_reason	❌ incomplete AttentionContract.provenance / AttentionBackendCapability	IDENTICAL (model shape, dtype, TP/CP) + SEMANTIC (mask, scale, RoPE/cache identity) + RECORDED (backend/fallback)	⚠️ probe-only; current reports are partial	#235 PR1 add required provenance/capability fields and fail closed on blocking mismatches	🙋
2	Materialization boundary	Full-prefill training graph vs. chunked-prefill / paged-prefill / decode rollout graph	✅ top-level attention; recorded field attention.materialization	❌ no first-class AttentionMaterializationSpec	SEMANTIC (same logical Q/K/V rows) + RECORDED (physical materialization)	⚠️ single-GPU attribution exists in #253/#260; distributed gate incomplete	#235 PR2/PR6 keep references; PR5 run full/chunked/paged/decode matrix	🙋
3	RoPE fusion and tensor state	Whether drift comes from fused vs. unfused RoPE, pre/post-RoPE Q/K, position ids, theta, or cast boundary	✅ top-level attention; recorded fields attention.rope_fusion_boundary, attention.q_rope_state, attention.k_cache_rope_state	❌ incomplete RoPE/cache fields in AttentionContract	SEMANTIC (position_ids, cache_position, theta, rotary layout, tensor state) + RECORDED (fusion)	⚠️ PR2/PR6 cover reference attribution; PR7 FlashInfer path is scaffolded	#235 PR1 make RoPE state first-class; PR7 validate TE/FlashInfer fused candidates	🙋
4	GQA head mapping and QKV layout	Whether drift comes from different head ownership, KV-head replication, packed/varlen layout, or transposes	✅ top-level attention; recorded fields attention.qkv_layout, attention.gqa_mapping, attention.tp_head_ownership	❌ incomplete AttentionLayoutSpec / TP head ownership	IDENTICAL (Hq/Hkv/D, TP partition) + SEMANTIC (GQA map) + RECORDED (layout)	⚠️ reference path exists, but capability checks are incomplete	#235 PR1 add layout/head-ownership fields; PR2/PR3 compare canonical layout against candidates	🙋
5	Paged KV and cache identity	Whether physical page order, prefix-cache reuse, slot mapping, or cache positions change the logical KV sequence	✅ top-level attention; recorded fields attention.paged_kv_policy, attention.cache_identity	❌ incomplete DecodeKVCacheMetadata in the main contract	SEMANTIC (logical token order, cache identity, position ids) + RECORDED (physical page table)	⚠️ single-GPU decode replay exists in #260; distributed/cache integration incomplete	#235 PR6 validate page/cache identity; PR7 bind FlashInfer paged attention behind those checks	🙋
6	LSE export/domain	Whether (out,lse) mismatch comes from missing LSE, natural-log vs. log2 LSE, wrong shape, or vocab-domain LSE	✅ top-level attention; recorded fields attention.lse_export, attention.lse_domain, attention.lse_base	❌ no complete AttentionLseSpec	SEMANTIC (lse domain/base/shape/masked-row behavior)	⚠️ reference and TE merge oracle exist; backend LSE capability is not fully gated	#235 PR1 add LSE spec; PR2/PR3/PR6/PR7 require/export compatible attention-domain LSE	🙋
7	CP / split-K reduction and merge order	Whether drift is caused by CP transport, split-K tiling, backend-local reduction order, or online-softmax merge order	✅ top-level attention; recorded fields attention.merge_order, attention.cp_comm_type, attention.split_kv_policy	❌ incomplete AttentionReductionSpec	SEMANTIC (global_block_index, acc dtype, merge order) + RECORDED (cp_comm_type, split policy, engine)	⚠️ deterministic CP reference exists in #238; split-K backend gates incomplete	#235 PR3 own fixed-order CP merge; PR5 run distributed split-K matrix; PR7 reject unsafe auto split-K under batch-invariant claim	🙋
8	Precision and downcast boundary	Whether drift comes from BF16/FP32 accumulation, backend internal dtype, partial-state downcast, or final write downcast	✅ top-level attention; recorded fields attention.accum_dtype, attention.downcast_at	❌ incomplete precision fields in AttentionReductionSpec	SEMANTIC (accum_dtype, downcast_at) + RECORDED (backend internal policy when observable)	⚠️ reference path can measure drift, but optimized backend provenance is partial	#235 PR1 add precision/downcast fields; PR5/PR7 report candidate precision behavior	🙋
9	TE / FlashInfer backend reuse	Whether Megatron-style TE training attention or vLLM-style FlashInfer rollout attention can satisfy the same RL-Kernel contract	✅ top-level attention; recorded fields attention.backend_candidate, attention.backend_capability	❌ no complete candidate capability spec	RECORDED (library/version/backend) + SEMANTIC only after capability checks pass	⚠️ TE merge oracle and FlashInfer scaffold exist; production promotion not complete	#235 PR7 evaluate TE full-prefill and FlashInfer paged prefill/decode candidates; no silent fallback	🙋
10	Training backward state	Whether forward-compatible attention also produces compatible dq/dk/dv under the training path	✅ top-level attention; recorded field attention.backward_backend	❌ no AttentionBackwardSpec / saved-state compatibility field	SEMANTIC (dq/dk/dv, saved forward state) + RECORDED (backend/autograd path)	❌ PR8 not implemented	#235 PR8 add backward reference and candidate oracle only if compatible saved state is available	🙋

Important boundary: Megatron/TE and vLLM/FlashInfer may use different physical kernels, layouts, pages, chunks, and communication transports. They may not disagree on Attention semantics that affect the result: mask, scale, GQA mapping, RoPE position/state, cache identity, TP/CP ownership, LSE convention, split-K policy, reduction order, and final downcast boundary.

LogProb knobs

Collected from @KJLdefeated 's study of how vLLM and Megatron implements logp kernel differently.

#	Axis	What it isolates	Knob	Contract field	Binding tier	Runnable today?	Task	Owner
1	Logprob runtime/provenance	Whether drift is arithmetic at all, or just two differently-configured processes	✅ top-level logprob; recorded logprob.impl (mcore_unfused / mcore_native_fused / mcore_te_fused / slime_vocab_parallel / vllm_all_gather), logprob.tp_size, logprob.requested_backend, logprob.actual_backend, logprob.fallback_reason	❌ no LogProbContract.provenance / LogProbBackendCapability	IDENTICAL (V, TP size, model dtype) + SEMANTIC (transform chain, mask) + RECORDED (impl/version)	⚠️ probe-only; nothing rejects a mismatch	#241 PR1 contract + dispatch metadata, fail closed on missing vocab-shard bounds	🙋
2	Vocab padding domain	Whether dlogp is a systematic bias from padded rows in the denominator rather than rounding	✅ logprob.vocab_padding_policy ∈ {include, mask}; recorded padded_vocab_size, org_vocab_size, num_pad_rows, pad_row_init	❌ no LogProbVocabSpec (padded vs. real vocab, mask policy)	SEMANTIC — the two sides must agree on which vocab entries are in the LSE	❌ MCore has no mask; vLLM has no include-mode	#241 PR1 make padded-vs-real vocab first-class; PR2 reference honors both modes; offline probe of lm_head[org_V:] norms	🙋
3	Vocab shard ownership	Whether partial-level disagreement is arithmetic or just different shard boundaries	✅ logprob.vocab_shard_map; recorded per-rank (vocab_start, vocab_end) for both engines	❌ no LogProbShardingSpec	IDENTICAL (global vocab domain) + RECORDED (per-rank boundaries)	❌ MCore VocabUtility and vLLM _get_indices disagree at TP≥2 for Qwen3	#241 PR1 shard map in metadata; PR4 cross-engine boundary reconciliation before partial comparison	🙋
4	Reduction topology	Whether drift comes from partial-LSE all-reduce vs. full-logit all-gather + single-pass softmax	✅ logprob.reduce ∈ {allreduce_partials, allgather_partials_fixed_merge, allgather_logits}; recorded collective, engine	❌ no LogProbReductionSpec	SEMANTIC (order, acc_dtype) + RECORDED (collective, engine)	❌ no fixed-order reference reduction exists	#241 PR2 the deterministic vocab-parallel reference; PR4 run all three arms	🙋
5	Merge order / NCCL as reducer	Whether NCCL algorithm choice, not the math, moves the number	✅ logprob.merge_order ∈ {nccl_algo, global_vocab_shard_index}; recorded NCCL_ALGO, NCCL_PROTO, world topology	❌ no order field	SEMANTIC (global_vocab_shard_index) + RECORDED (NCCL algo)	⚠️ TE fused already pins order; MCore unfused/native do not	#241 PR2 own the fixed-order merge; PR4 repeat-run bitwise assertion under varied NCCL_ALGO	🙋
6	Precision and downcast boundary	Whether drift comes from fp32 upcast position, head GEMM output dtype, or partial downcast	✅ logprob.head_dtype, logprob.accum_dtype, logprob.downcast_at	❌ no precision fields	SEMANTIC (accum_dtype, downcast_at) + IDENTICAL (head GEMM output dtype)	⚠️ selectable on both sides, not declarable	#241 PR1 add precision fields; PR4 sweep head_dtype ∈ {bf16, fp32} against an fp32 oracle	🙋
7	Logit transform chain and temperature	Whether the two sides are computing the same distribution at all	✅ logprob.transform_chain (ordered list); recorded temperature_applied, logprobs_mode, soft_cap, logit_scale	❌ no LogProbTransformSpec	SEMANTIC — an ordering mismatch here is a correctness bug, not drift	⚠️ both sides configurable, neither declares	#241 PR1 make the transform chain first-class and fail closed on mismatch	🙋
8	Selected-logit gather and token alignment	Whether drift comes from single-owner gather, ignore_index, OOV targets, or a shift-by-one	✅ logprob.selected_gather ∈ {masked_sum_allreduce, local_gather}; recorded ignore_index, active_mask, shift_convention	❌ no single-owner / mask fields	SEMANTIC (single-owner invariant, pre-shifted input convention)	⚠️ MCore masked-sum vs. vLLM local gather; no cross-check	#241 PR2 masked single-owner gather + exact-tiling validation	🙋
9	LM-head GEMM and materialization boundary	Whether the number moved before any softmax ran — full-sequence training GEMM vs. logits_indices decode GEMM with CUDA-graph batch padding	✅ logprob.head_gemm_kernel, logprob.materialization ∈ {full_prefill, chunked, last_token_only}; recorded padded batch size	❌ no LogProbMaterializationSpec	IDENTICAL (logical rows) + RECORDED (physical M, kernel, tile/split-K)	❌ inherits #249 GEMM row 2; no shard reference at the head	#239 PR2 head GEMM reference; #241 PR4 report head-GEMM drift separately from LSE drift	🙋
10	Fusion tiling	Whether TE's BLOCK_SIZE online-softmax tiling, native jit-fusion, or slime's sequence chunking changes the partial	✅ logprob.fusion_impl ∈ {none, native, te, chunked_linear_ce}; recorded block_size, chunk_size	❌ no fusion field	RECORDED (measured, never forced equal)	⚠️ all three exist upstream; none dispatchable through the catalog	#241 PR4 run the fusion matrix against PR2's reference	🙋
11	Training-only state (grad, entropy)	Whether a forward-compatible logprob also produces a compatible dlogits and entropy	✅ logprob.backward_impl, logprob.grad_downcast_at, logprob.entropy_reduce	❌ no LogProbBackwardSpec	SEMANTIC (dlogits, entropy reduction) + RECORDED (autograd path)	❌ not implemented	follow-up issue: vocab-parallel logprob backward + GRPO loss (blocked on #241)	🙋