Repository navigation
[WS1][CUDA][MiniMax-H3] projection rows: audio in/out, text context and QKV on a fixed-order GEMM (RFC #420 stack 2/5) - #568
Conversation
|
Important Review skippedToo many files! This PR contains 132 files, which is 32 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.
 |
|
I was planning to work on the audio projections, as I’m currently developing the video input/output projections. After seeing your PR, I briefly reviewed the implementation and noticed that I reproduced this with the unmodified operator at import torch
from rl_engine.backends.cuda.model_specific.minimax_h3.projection import H3AudioInputProjectionCudaOp
op = H3AudioInputProjectionCudaOp()
x = torch.ones(3, 32, device="cuda")
w = torch.zeros(5376, 32, device="cuda", requires_grad=True)
b = torch.zeros(5376, device="cuda", requires_grad=True)
g = torch.zeros(3, 5376, device="cuda")
g[:, 0] = torch.tensor([1.0, 2**-24, -1.0], device="cuda")
def grads(order):
return torch.autograd.grad(op(x[order], w, b), (w, b), g[order])
dw1, db1 = grads([0, 1, 2])
dw2, db2 = grads([0, 2, 1])
print("dW[0,0]:", dw1[0, 0].item(), dw2[0, 0].item())
print("dbias[0]:", db1[0].item(), db2[0].item())
assert torch.equal(dw1, dw2) and torch.equal(db1, db2), "parameter-gradient permutation BI failed"Output: The exact value for both printed components is |
…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>
Hi @walker-ai, first of all, I'm sorry for taking the audio projection rows while you were already working on them. I should have checked with you before claiming and opening the PR. Thank you for the detailed review and the reproducer. You were right: dW and dbias were reduced in call order, so jointly permuting x and g changed them. I've followed your suggestion and the canonical-row design from #562 in 0da0583 and d2941f0: all projection ops (audio in/out, text context, Q/K/V) take an optional logical_keys ([T, 2|3] int64, negative first column = padding); rows are reduced in ascending key order with padding dropped, and active keys must be unique; |
…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>
0da0583 to
3cf51bf
Compare
…gn#568 Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…#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>
Thanks for the fix and the co-author credit! I'd be happy to help. As a first step, I can adapt our BI checks to your updated FP32 audio projections, focusing on whole-batch versus microbatch execution, row ordering, and padding. Once we have the results, we can agree on any implementation changes and how to split them. The existing canonical backward session used by #562 could be a reuse option if cross-call accumulation needs changes. |
Flink-ddd
added
platform: cuda
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
What
Four linear rows on one fixed-order Triton GEMM, contract
h3-projection-gemm-v1:BLOCK_Ksteps into FP32 registers;no split-K, stream-K, atomics or autotuning; the tile shape does not depend on
M,Nor the grid;tl.dot, tile 128x128x64, 8 warps, 3 stages; FP32 operands:input_precision="ieee"(no TF32), tile 64x64x32, 4 warps, 3 stages;dX: reduction over output features;dW: reduction over tokens ascending;db: ascending token sum).audio_input_projection_fp32audio_output_projection_fp32text_context_projectionh3_qkv_gemm[Q;K;V]form equals the separate ones bit for bit)Prior art & reuse decision
F.linearM(split-K), so it is not batch-invariant and cannot be made fail-closed; kept as the provider baselineQKVParallelLinearforward_fusedh3-det-linear-v1(this repo)Decision: a new fixed-order Triton GEMM.
Results (H100 / Hopper)
Measured on NVIDIA H100 (93 GiB, sm90), driver 580, CUDA 13.0, torch 2.13.0+cu130, pinned
MiniMax-H3@42ed227weights, at head3cf51bf9. Relative error is against an FP64F.linear; time is the median forward time.audio_input_projection_fp32audio_input_projection_fp32audio_input_projection_fp32audio_output_projection_fp32audio_output_projection_fp32audio_output_projection_fp32text_context_projectiontext_context_projectiontext_context_projectionh3_qkv_gemmh3_qkv_gemmh3_qkv_gemmNumbers:
bench.json. The kernel is 1.3x-4.3x slower than cuBLAS on the BF16 rows and3.5x-12x slower on
audio_output_projection_fp32(K = 5376, FP32 ieee, 64x64x32 tile), whose error is also about 3x the provider's (1.3e-6 vs 4.4e-7; accuracy floor relaxed to 4e-6, see the doc).That row is the first candidate for a per-row tile sweep. This is the price of the fixed reduction order; strict mode forbids split-K.
Results (B200 / Blackwell)
Measured on NVIDIA B200 (179 GiB, sm100), driver 580.126, CUDA 13.0, torch 2.13.0+cu130, pinned
MiniMax-H3@42ed227weights, clean tree at3cf51bf9.Tests (
test_h3_projection_rows.pyandtest_h3_one_block.py): 67 passed (67 collected);tests.json.Forward time and accuracy against FP64 for each row against the provider's linear (one B200, 20 repeats;
bench.json):audio_input_projection_fp32audio_input_projection_fp32audio_input_projection_fp32audio_output_projection_fp32audio_output_projection_fp32audio_output_projection_fp32text_context_projectiontext_context_projectiontext_context_projectionh3_qkv_gemmh3_qkv_gemmh3_qkv_gemmThe fixed-order kernels are slower than the provider's GEMM (ratio above 1 in every row, up to 8.33x) and match it in accuracy, except
audio_output_projection_fp32, where the candidate's error is 1.3e-06 against 4.3e-07 for the provider at 64 tokens (both far below the BF16 level of the other rows).Status
Forward and backward are bitwise repeatable and invariant to batch size, position, permutation and padding. The
audio_outputrow has no legal output shard (see Scope).Tests
python -m pytest tests/models/minimax_h3/test_h3_projection_rows.py tests/models/minimax_h3/test_h3_one_block.py -q # 61 passed, 0 skippedCovers real checkpoint tensors (shard 1: audio, context embedder; block 0: Q/K/V), FP64 accuracy, repeat and backward determinism, fused vs separate QKV, TP output shards, and fail-closed rejections (wrong widths, mixed devices, CPU tensors, wrong dtype).
B200: same files, 67 passed.
Scope / Known limitations
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.
audio_output_projection_fp32has a 32-wide output equal to one shard unit, so it has no legal output shard; a test asserts that narrower weights are rejected.Not part of this PR: video projections (#562), O / gate_up / down GEMMs (#554, #535), full attention (#553), ROCm, SGLang adapter, pipeline parallel.