Skip to content

[WS1][CUDA][MiniMax-H3] token_refiner_block: two refiner blocks over the strict operator rows (RFC #420 stack 4/5) - #570

Open
fusheng-ji wants to merge 19 commits into
RL-Align:test-h3from
fusheng-ji:feat/h3-token-refiner
Open

fusheng-ji wants to merge 19 commits into
RL-Align:test-h3from
fusheng-ji:feat/h3-token-refiner

Conversation

@fusheng-ji

@fusheng-ji fusheng-ji commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

RFC #420 H3 stack 4/5 (base test-h3, merge in this order): #586 -> #587 -> #588 -> foundation -> projection rows -> WS2 parallel -> token refiner -> full transformer. This PR merges after its parent; until then GitHub also shows the parent's commits. This PR's own diff.

Evidence location: the measurement outputs of this PR (JSON, figures) are not part of its diff. They are on the fork's pr_assets branch under pending/h3-token-refiner/ (links in this description are fixed to that commit). Paths such as reports/experiments/... below name where a runner writes its output.

Acceptance evidence (maintainer): Hopper H100 (sm90) on pr_assets @ d188d29e. Blackwell (B200) below is supplementary and is not a Hopper merge blocker.

Depends on the foundation PR (the branch contains it as a merge) and the projection rows.

What

  • The text-stream refiner (token_refiner.blocks.0/1, 14 nodes each, no AdaLN, no RoPE, then final_norm) as a node graph in h3_refiner.py, with forward and backward,
    provider and FP64 golden replays, and the same per-node comparison harness as the transformer block.
  • The row owns the two residual adds (FP32 add, one rounding). Every other node is bound to an existing row (h3_rmsnorm, h3_qkv_gemm) or stays interim for rows
    other contributors own (attention, SwiGLU, O / gate_up / down).
  • Operator trace gets a token_refiner stage; h3_weights.load_h3_tensors reads real tensors from the sharded checkpoint.
  • The node vocabulary equals fingerprint.REFINER_NODES (checked by a test).

Prior art & reuse decision

diffusers MiniMaxH3TokenRefiner is the reference graph; vLLM-Omni reuses its generic block, so there is no separate fused implementation to adopt. No new kernel is written here.
Decision: compose existing rows; do not re-implement attention or the FFN GEMMs.

Results (H100 / Hopper)

Hopper pytest summary: tests.json (132 collected, 132 passed @ 6de320dc, H100 sm90).

Hopper test gate (PR head 6de320dc): test_h3_token_refiner.py and the stacked files in Tests below — 132/132 pass on H100 per tests.json (including test_nodes_accuracy_repeat_batch_and_isolated[1-257] / swiglu). The six-case chained bound matrix is swiglu_tolerance_matrix.json (candidate ~0.0121 at block 1 / 257 tokens; provider chain matches on H100).

Measured on NVIDIA H100 (93 GiB, sm90), driver 580, CUDA 13.0, torch 2.13.0+cu130, pinned MiniMax-H3@42ed227 weights, at head 17d0e9e9 (forward path unchanged, numbers not re-measured). Real block-0 weights, one block, forward.

Tokens RL-Kernel graph Provider Ratio Output error vs FP64 (RL-Kernel) (provider)
12 3.68 ms 0.54 ms 6.8x 2.92e-03 2.92e-03
257 8.49 ms 0.62 ms 13.7x 5.72e-03 5.72e-03
925 32.28 ms 1.42 ms 22.8x 6.65e-03 6.65e-03

refiner accuracy and performance

Numbers: report.json. Repeat and batch invariance hold at every node. The first node that differs from the
provider is the attention node. The time gap is almost entirely the interim bindings of rows other contributors own; the nodes this PR and the projection rows add are within about 2x of the provider.

Results (B200 / Blackwell) — supplementary, not Hopper acceptance

Measured on NVIDIA B200 (179 GiB, sm100), driver 580.126, CUDA 13.0, torch 2.13.0+cu130, pinned MiniMax-H3@42ed227 weights, clean tree at 6de320dc.

One test fails on B200. test_h3_token_refiner.py::test_nodes_accuracy_repeat_batch_and_isolated[1-257] (block 1, 257 tokens), at the swiglu node: the candidate chain's relative error against the FP64 golden is 0.01211 and the bound is 1.25 x provider + 1e-3 = 0.01116, with the provider at 0.00816. Result of the test file: 131 passed, 1 failed (132 collected); tests.json. The bound was not changed. On H100 the same case passes (candidate and provider chains both 0.012117; the six-case matrix is swiglu_tolerance_matrix.json).

What a diagnostic on the B200 showed (separate runs, not part of the result above):

  • The candidate's swiglu error is the same on both devices (0.01211 on B200, 0.012117 on H100); what differs is the provider chain (0.00816 on B200, 0.012117 on H100), which lowers the bound.
  • The swiglu kernel is not less accurate than the provider: given the same input its error equals the provider's (0.00816), and on random BF16 input against FP64 it is more accurate (0.0017 against 0.0024) because it rounds once after FP32 math where the provider rounds twice.
  • The test compares chained graphs, so a node's error includes everything upstream. In this case ffn_gate_up is equal for both chains (0.00818) and swiglu is 0.01211 (candidate) against 0.00816 (provider). The direction is not systematic: block 0 with 257 tokens has the candidate below the provider (0.00829 against 0.01081) on the same device.

Forward time of one refiner block (report.json):

tokens candidate ms provider ms ratio
12 2.19 0.23 9.4x
257 4.88 0.34 14.3x
925 23.37 0.69 33.7x

token refiner on B200: accuracy per node and forward time

Repeat and batch invariance hold in every case (first_repeat_drift and first_batch_drift are none); the first difference from the provider is the interim attention node, as on H100.

Status

Hopper (merge gate): graph vocabulary pinned to the fingerprint; forward and backward on real weights; repeat/batch invariance on H100; refiner pytest 132/132 pass on H100 at head 6de320dc.

B200 (supplementary): 131/132 pytest pass; one chained-graph swiglu bound fail (provider chain lower than on H100, not a worse candidate kernel — see swiglu_tolerance_matrix.json). No bound was widened; not treated as blocking Hopper review.

Tests

python -m pytest tests/models/minimax_h3/test_h3_token_refiner.py tests/models/minimax_h3/test_h3_operator_trace.py \
    tests/models/minimax_h3/test_h3_arch_fingerprint.py tests/models/minimax_h3/test_h3_projection_rows.py \
    tests/models/minimax_h3/test_h3_one_block.py -q                                                    # 126 passed

B200: same files, 131 passed, 1 failed (see Results (B200 / Blackwell)).

Scope / Known limitations

Evidence priority: maintainers reviewing RFC #420 on CUDA should use H100 sections and JSON under pr_assets@d188d29e. B200 numbers document sm100 behaviour only.

Strict mode is fail-closed: no TF32, no fast-math, no split-K; FP32 accumulation with one final rounding; unsupported shapes, dtypes and devices raise instead of falling back. Not part of this PR: video projections (#562), O / gate_up / down GEMMs (#554, #535), full attention (#553), ROCm, SGLang adapter, pipeline parallel. Replacing the interim bindings is tracked by those rows.

@coderabbitai

coderabbitai Bot commented Oct 22, 2026 •

Copy link
Copy Markdown

Important

Review skipped

Too many files!

This PR contains 151 files, which is 51 over the limit of 100.

To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch.

Upgrade to a paid plan to raise the limit.

This review couldn't start because sufficient usage credits or metered capacity aren't available. Add credits or update usage-based reviews in the billing tab, then retry.

Check out review usage here.

Run configuration
Commits

Reviewing files that changed from the base of the PR and between d22d4b4 and 3803588.

Files selected for processing: 151
  • .github/workflows/ci.yml
  • benchmarks/models/benchmark_h3_conditioning.py
  • build_tools/extensions.py
  • csrc/bindings/ops.cpp
  • csrc/cuda/h3/adaln_row_gather.cu
  • csrc/cuda/h3/det_linear.cu
  • csrc/cuda/h3/gate_residual.cu
  • csrc/cuda/h3/rmsnorm_modulate.cu
  • csrc/cuda/h3/timestep_sinusoid.cu
  • docs/.nav.yml
  • docs/operators/README.md
  • docs/operators/h3-adaln-gate-residual.md
  • docs/operators/h3-adaln-projection.md
  • docs/operators/h3-adaln-row-gather.md
  • docs/operators/h3-arch-fingerprint.md
  • docs/operators/h3-final-adaln-out.md
  • docs/operators/h3-one-block.md
  • docs/operators/h3-operator-trace.md
  • docs/operators/h3-projection-rows.md
  • docs/operators/h3-rmsnorm.md
  • docs/operators/h3-sp-norm-adaln.md
  • docs/operators/h3-timestep-mlp.md
  • docs/operators/h3-timestep-sinusoid.md
  • docs/operators/h3-token-refiner.md
  • docs/operators/h3-tp-adaln-3mod.md
  • docs/operators/h3-trajectory-artifact.md
  • examples/models/minimax_h3/README.md
  • examples/models/minimax_h3/candidate.sbatch
  • examples/models/minimax_h3/diagnose-block.sbatch
  • examples/models/minimax_h3/full-model-ensemble-rule.json
  • examples/models/minimax_h3/full-model-thresholds.json
  • examples/models/minimax_h3/generation-compare.sbatch
  • examples/models/minimax_h3/generation-contract.json
  • examples/models/minimax_h3/generation-inputs.json
  • examples/models/minimax_h3/generation-prompt.txt
  • examples/models/minimax_h3/generation-resume-check.sbatch
  • examples/models/minimax_h3/generation.sbatch
  • examples/models/minimax_h3/reference-manifest.json
  • examples/models/minimax_h3/reference.sbatch
  • rl_engine/_C.pyi
  • rl_engine/backends/cuda/model_specific/__init__.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/projection.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py
  • rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py
  • rl_engine/backends/shared/triton/gemm/h3_projection.py
  • rl_engine/models/minimax_h3/arch_fingerprint.json
  • rl_engine/models/minimax_h3/fingerprint.py
  • rl_engine/models/minimax_h3/operator_trace.py
  • rl_engine/models/minimax_h3/pinned_checkpoint.json
  • rl_engine/models/minimax_h3/trajectory.py
  • rl_engine/reference/minimax_h3/__init__.py
  • rl_engine/reference/minimax_h3/adaln_projection.py
  • rl_engine/reference/minimax_h3/adaln_row_gather.py
  • rl_engine/reference/minimax_h3/block.py
  • rl_engine/reference/minimax_h3/final_adaln_out.py
  • rl_engine/reference/minimax_h3/fixed_order.py
  • rl_engine/reference/minimax_h3/gate_residual.py
  • rl_engine/reference/minimax_h3/projection.py
  • rl_engine/reference/minimax_h3/rmsnorm.py
  • rl_engine/reference/minimax_h3/timestep_mlp.py
  • rl_engine/reference/minimax_h3/timestep_sinusoid.py
  • rl_engine/runtime/registry.py
  • rl_engine/validation/models/h3_block.py
  • rl_engine/validation/models/h3_cases.py
  • rl_engine/validation/models/h3_chain.py
  • rl_engine/validation/models/h3_full_manifest.json
  • rl_engine/validation/models/h3_full_model.py
  • rl_engine/validation/models/h3_generation_resume.py
  • rl_engine/validation/models/h3_manifest.json
  • rl_engine/validation/models/h3_provider.py
  • rl_engine/validation/models/h3_refiner.py
  • rl_engine/validation/models/h3_report.py
  • rl_engine/validation/models/h3_weights.py
  • rl_engine/validation/models/h3_ws2.py
  • rl_engine/validation/operators/operator_inputs.py
  • rl_engine/validation/operators/operator_specs.py
  • tests/build/test_build_platform_collectives.py
  • tests/models/minimax_h3/conftest.py
  • tests/models/minimax_h3/test_h3_adaln_gate_residual.py
  • tests/models/minimax_h3/test_h3_adaln_modulation.py
  • tests/models/minimax_h3/test_h3_adaln_projection.py
  • tests/models/minimax_h3/test_h3_adaln_row_gather.py
  • tests/models/minimax_h3/test_h3_arch_fingerprint.py
  • tests/models/minimax_h3/test_h3_benchmark_timing.py
  • tests/models/minimax_h3/test_h3_chain_replay_cli.py
  • tests/models/minimax_h3/test_h3_cli.py
  • tests/models/minimax_h3/test_h3_conditioning_e2e.py
  • tests/models/minimax_h3/test_h3_det_linear.py
  • tests/models/minimax_h3/test_h3_final_adaln_out.py
  • tests/models/minimax_h3/test_h3_full_model.py
  • tests/models/minimax_h3/test_h3_generation_resume.py
  • tests/models/minimax_h3/test_h3_native_validation.py
  • tests/models/minimax_h3/test_h3_one_block.py
  • tests/models/minimax_h3/test_h3_operator_trace.py
  • tests/models/minimax_h3/test_h3_projection_rows.py
  • tests/models/minimax_h3/test_h3_report.py
  • tests/models/minimax_h3/test_h3_rmsnorm.py
  • tests/models/minimax_h3/test_h3_sp_norm_adaln.py
  • tests/models/minimax_h3/test_h3_timestep_mlp.py
  • tests/models/minimax_h3/test_h3_timestep_sinusoid.py
  • tests/models/minimax_h3/test_h3_token_refiner.py
  • tests/models/minimax_h3/test_h3_tp_adaln.py
  • tests/models/minimax_h3/test_h3_trajectory_artifact.py
  • tests/models/minimax_h3/test_h3_weights.py
  • tests/models/minimax_h3/test_h3_ws2_processes.py
  • tests/models/minimax_h3/test_prepare_h3_weights.py
  • tests/runtime/test_dispatch.py
  • tests/validation/operators/test_ws1_gtest_gpu.py
  • tools/migration/README.md
  • tools/migration/layout.json
  • tools/validation/models/h3_bitwise_reuse.py
  • tools/validation/models/h3_block_diagnostic.py
  • tools/validation/models/h3_block_replay.py
  • tools/validation/models/h3_block_ulp_sweep.py
  • tools/validation/models/h3_candidate_prepare.py
  • tools/validation/models/h3_chain_replay.py
  • tools/validation/models/h3_evidence.py
  • tools/validation/models/h3_evidence_integrity.py
  • tools/validation/models/h3_fingerprint.py
  • tools/validation/models/h3_first_step_diagnostic.py
  • tools/validation/models/h3_full_compare.py
  • tools/validation/models/h3_full_ensemble_rule.py
  • tools/validation/models/h3_full_generation.py
  • tools/validation/models/h3_full_model.py
  • tools/validation/models/h3_full_provider_spread.py
  • tools/validation/models/h3_generation_adopt.py
  • tools/validation/models/h3_generation_compare.py
  • tools/validation/models/h3_operator_trace.py
  • tools/validation/models/h3_prior_art.py
  • tools/validation/models/h3_projection_bench.py
  • tools/validation/models/h3_projection_charts.py
  • tools/validation/models/h3_refiner_report.py
  • tools/validation/models/h3_trajectory.py
  • tools/validation/models/h3_ws2_evidence.py
  • tools/validation/models/plot_h3_bitwise_reuse.py
  • tools/validation/models/plot_h3_block.py
  • tools/validation/models/plot_h3_evidence.py
  • tools/validation/models/plot_h3_prior_art.py
  • tools/weights/prepare_h3_weights.py
  • tools/weights/snapshot_h3_checkpoint.py

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

fusheng-ji and others added 16 commits October 22, 2026 22:54
…fixed-order Triton GEMM

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 650e2148c6f0e5dcdef136537e0e62f745020273)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit d099c0b)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b41f868)
…il-closed test instead of a skip

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 74e5b83)
…l keys

Review finding (RL-Align#568): jointly permuting x and g changed dW/dbias because the token axis
was reduced in call order. All projection ops (audio in/out, text context, Q/K/V) now take an
optional logical_keys [T,2|3] (negative first column = padding); rows are reduced in ascending
key order with padding dropped, active keys must be unique. Key format follows RL-Align#562.
Adds the reviewer's repro plus permutation/padding tests for every row and Q/K/V.

Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit d2941f0)
…ready ordered rows

Follows the canonical-row design of RL-Align#562: the Q/K/V node of the block graph supplies
(batch row, position) keys so its dW is reduced in logical order, and rows that already arrive in
canonical order are reduced in place (same bytes, no gather of x and g).

Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 0da0583)
…block 0 node-by-node and 50-block/generation gates

Squashes the one-block branch (its row-by-row history: timestep sinusoid, RMSNorm, AdaLN rows,
sequence/tensor parallel norm rows, final AdaLN, one block, 50-block qualification) onto the current
test-h3 so the pull request carries a reviewable, signed history. RFC RL-Align#420.

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…pull request

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit eaba5efe3df971ea9d1a969f53e48f09e09ed176)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 7da0574)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b45190d08fc033c7e63ad821d3918eb32dd454f6)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b5536d6)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit ad1925045b522772bf56f8468b23cca47f04156b)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit f9041e8)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b44540e)
…ody)

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit eab901b)
… weight loader, tests, doc

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit d12a6bb26566fae58a309e76916af9aae767d487)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b316869)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 1e3a2ea)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 67af796)
…age import

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
@fusheng-ji
fusheng-ji force-pushed the feat/h3-token-refiner branch from 67af796 to 6de320d Compare October 22, 2026 21:04
fusheng-ji added a commit to fusheng-ji/RL-Kernel that referenced this pull request Oct 22, 2026
…#570

Add a sanitized JSON replay of all six accuracy-test swiglu nodes on Hopper,
with a cross-device reference to the B200 block=1/257 bound failure.
fusheng-ji added a commit to fusheng-ji/RL-Kernel that referenced this pull request Oct 23, 2026
…#567 RL-Align#568 RL-Align#570 RL-Align#571 RL-Align#574 RL-Align#575

Test summaries (tests.json), runner outputs and figures measured on NVIDIA B200 (sm100),
driver 580, CUDA 13.0, torch 2.13.0+cu130, MiniMax-H3@42ed227, clean trees at the PR heads
eab901b, 3cf51bf, 6de320d, eb7840b, af3dcae and 4f7c5d1. Figures use B200 axis labels.
No host names, paths or scheduler identifiers.

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
fusheng-ji added a commit to fusheng-ji/RL-Kernel that referenced this pull request Oct 23, 2026
…ory step_max on H100

- h3-token-refiner-h100/tests.json (132/132 @ 6de320d)
- h3-full-transformer-h100/tests.json (142/142 @ eb7840b)
- h3-full-transformer-h100/dual_trajectory_step_max.json (auditable per-step max)

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
@Flink-ddd Flink-ddd added platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations) new-model MiniMax-H3 WS1 labels Oct 23, 2026
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

MiniMax-H3 new-model platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations) WS1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants