Feat/nemotron h lm head logprob cuda

[effintell]

Implements the CUDA column of the RFC #434 row "Embedding lookup (untied), LM head 2688→131072 fused selected-token logprob (no softcap)" for Nemotron-3-Nano (extends #243, #204, #336).

• NemotronHLMHeadLogprobSM90Op: strict, batch-invariant fused LM-head selected-token logprob (forward never materialises the [T, 131072] logits) + optional per-row temperature.
• Strict backward: deterministic dlogits kernel + det-GEMM for grad_hidden / grad_weight, chunked by a fixed element budget (chunking is bit-neutral).
• Untied embedding lookup with a sorted-segment deterministic grad_weight.
• Fail closed: geometry gate (hidden 2688 / vocab 131072 / BF16 / SM90) and strict det-GEMM backend requirement raise instead of silently falling back; provenance (contract / checkpoint / kernel ids / det-GEMM backend id / fallback flag) is recorded and checked.
• Fix: the SM90 det-GEMM tree reduces at most K = 32768 (TREE_DEPTH=10). Backward reductions with K > 32768 (grad_hidden over vocab = 131072; grad_weight over tokens > 32768) previously crashed with an illegal memory access. They are now split into ≤ 32768 segments whose FP32 partials are summed in a fixed ascending order (SM90 backend only; cuBLASLt no-split-K needs no split).
• grad_weight is judged against FP64 with the WS1 reduction gradient_accuracy tolerance (atol 0.1, rtol 0.02) instead of the logprob cell (atol 0.05, rtol 0) — see "Tolerance" below.

Out of scope: ROCm, TP / CP / EP.

Validation (H100 PCIe, torch 2.8.0+cu128, nvcc 12.8)

Build: KERNEL_ALIGN_FORCE_SM90=1 KERNEL_ALIGN_DET_GEMM_SM90=1 pip install -e . --no-build-isolation

#	Command	Result
2	python -m pytest tests/models/nemotron_h/test_lm_head_logprob_cuda.py::TestGeometryGate tests/runtime/test_dispatch.py -q (CPU)	24 passed
3a	python -m pytest tests/models/nemotron_h -q -rs — SM90 det-GEMM (auto, no cuBLAS env)	34 passed, 1 xfailed, 0 skipped
3b	same, cuBLASLt (CUBLAS_WORKSPACE_CONFIG=:16:8 CUBLASLT_WORKSPACE_SIZE=1)	31 passed, 3 skipped (SM90-only split tests, see below), 1 xfailed
4	`python tools/validation/operators/check_operator.py --op nemotron_h_lm_head_logprob --candidate cuda-sm90 --device cuda --dtype bf16 --batch 2 --seq 16 --vocab 131072 --normalized-dim 2688
--check-grad`	fails — checker issue, not this op: forward max_abs 0.67 vs atol 0.05; the existing linear_logp suite fails with the identical error on the same shape for both cuda-sm90 and triton
candidates
5	benchmarks/operators/logprob/benchmark_nemotron_h_lm_head_logp.py --json ...	see Benchmark
6	python -m pytest tests -q -k "linear_logp or fused_logp" -rs --ignore=tests/backends/rocm/test_rocm_aiter_api_contract.py	SM90: 101 passed, 30 skipped; cuBLASLt: 101 passed, 30 skipped (skips:
Ascend / MUSA / 2-GPU only)

The xfail is test_embedding_grad_weight_matches_fp64 (strict xfail): SM90EmbeddingOp accumulates repeated-id gradients in BF16; bitwise reproducibility of that gradient is tested and passes.
test_rocm_aiter_api_contract.py fails to collect on the base branch as well.

Acceptance checks

• Batch invariance (torch.equal): forward logprob / lse rows across batch sizes 1 / 2 / 7 / 64 / 255 / 257 (random row subsets), leading shape [3, 100], repeat runs, and train vs rollout entry;
  grad_hidden rows across batch sizes 1 / 7 / 64 / 257.
• FP64 parity: forward and both gradients allclose to an independent FP64 reference / VJP.
• Fail closed: geometry gate + det-GEMM backend gate raise; no silent fallback (provenance fallback: false).
• Provenance: recorded per run (det_gemm_backend_id, kernel ids, checkpoint nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16@bf77c31).

det-GEMM split verification (SM90 backend, profiler-confirmed sm90 kernels, rlkernel.det_gemm.sm90.v1)

All cases bitwise reproducible across two runs. "budget" = max_abs / allowed tolerance.

case	segments K	max_abs vs FP64	BF16 floor	reduction budget	logprob budget
dW rows=32768	[32768]	0.055	0.022	54%	109%
dW rows=32769	[32768, 1]	0.064	0.016	45%	128%
dW rows=32864	[32768, 96]	0.064	0.021	44%	129%
dW rows=65536	[32768, 32768]	0.090	0.031	65%	181%
dW rows=65600	[32768, 32768, 64]	0.089	0.031	58%	178%
dX V=32832	[32768, 64]	0.0013	0.0005	1%	3%
dX V=131072	4 × 32768	0.0022	0.0010	2%	4%

• Negative control: forcing the segment limit to 2^40 (no split) reproduces CUDA error: an illegal memory access for dW 32864, dW 65536 and dX 131072.
• test_grad_weight_token_reduction_beyond_det_gemm_tree is parametrised at exactly 32768 (1 segment), 32769 (K=1 tail) and 65600 (3 segments), asserts the segment count, and skips on non-SM90 backends.
  Disabling the split makes the 2- and 3-segment cases fail; dropping the last partial makes the 3-segment case fail (max diff 0.38 > 0.1).

Tolerance (grad_weight)

grad_weight is a reduction over tokens with entries up to |ref| ≈ 18. At that magnitude BF16 output rounding alone is 0.049 (half-ulp in [16, 32)), which exhausts the logprob cell's absolute-only budget
(atol 0.05). Measured on the 64-token model test: kernel max_err 0.051 vs BF16-rounding floor 0.049 (2 / 352M elements > 0.05). So it is judged with the reduction cell (atol 0.1, rtol 0.02); grad_hidden
stays on the logprob cell (max_err 0.0025).

Note: the SM90 det-GEMM tree accumulates in BF16 per level, so grad_weight error grows with token count — 42–54% of the reduction budget at ~32k tokens, 58–65% at ~65k tokens (≈ 2.5–4× the BF16 floor). The
split itself adds no measurable error.

Benchmark (Nemotron-3-Nano shape, H100 PCIe)

Train fwd+bwd, strict vs native PyTorch:

tokens	cuBLASLt strict / native (ms)	SM90 det-GEMM strict / native (ms)	peak MB native / strict (cuBLASLt)
2048	33.6 / 16.0 (2.1×)	109.7 / 16.0 (6.8×)	3072 / 1195
4096	85.7 / 32.3 (2.7×)	238.1 / 32.5 (7.3×)	6144 / 1717
8192	169.9 / 66.0 (2.6×)	480.5 / 66.3 (7.2×)	12288 / 1738
16384	351.2 / 140.0 (2.5×)	996.9 / 139.6 (7.1×)	24576 / 1780

Forward only (prefill): 2048 tokens 19.1 vs 6.4 ms (3.0×), 8192 tokens 64.7 vs 24.8 ms (2.6×), peak memory halved. cuBLASLt no-split-K is the recommended det-GEMM backend for training.

Known issues / follow-ups

• csrc/cuda/gemm/det_gemm_kernel.cu: when K is not a multiple of 32 or exceeds 32768, the SM90 det-GEMM silently falls back to the naive kernel instead of failing closed. This PR avoids the K > 32768 case
  from Python; the C++ behaviour will be tracked in a separate issue.
• The split tests only run under the SM90 det-GEMM build; a cuBLASLt-only CI skips them.
• Embedding grad_weight FP64 parity (BF16 accumulation) — strict xfail.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: cec28ba5-ab3f-41ad-8c26-285a4077ed22

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

Use the checkbox below for a quick retry:

• 🔍 Trigger review

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

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