feat(rocm): class-axis loss + fpquant/nvfp4 lanes — completes the ROCm loss/quant mirror - #172
Merged
Merged
Conversation
…m loss/quant mirror Final ROCm-mirror PR for the loss/quant families — closes the cross-device pair for the S11 class-axis losses and the S9 quantize ops (x86 in #167/#168/#169). class-axis loss (rocm_class_loss_compiled, pure composition — no new C++): cross_entropy / kl / js / focal / label_smoothed_cross_entropy / z_loss; exp/log run on the rocm unary lane (gfx1151), class-axis max/sum/gather/one-hot on the host (same pattern as the x86 class-loss + ROCm attention/gemm-family lanes). fpquant + nvfp4 (new GenerateROCMFpQuantKernel.cpp): tessera_rocm.fpquant — flat 1-operand float-grid mantissa-snap (ax=min(|x|,max); e=max(floor(log2 ax), min_exp); ulp=2^(e-mant); roundeven; sign+NaN-propagate), log2/floor/exp2/ roundeven via math->ROCDL, parameterized for fp8 e4m3/e5m2, fp6 e2m3/e3m2, fp4 e2m1. rocm_fpquant_compiled (per-tensor scale) + rocm_nvfp4_compiled (per-block fp8-E4M3 scale + E2M1 codes, host block structure). Validation (gfx1151): - tessera-opt codegen + ROCDL lowering verified. - test_rocm_class_loss_compiled.py + test_rocm_fpquant_compiled.py — 15 passed vs tessera.losses/tessera.ops (class-loss 2e-4; fpquant/nvfp4 2e-3 — the GPU log2/roundeven grid matches the reference on random data). Wiring: runtime executors (class-loss / fpquant / nvfp4) + `_rocm_fpgrid` / `_rocm_unary_t` / `_rocm_log_softmax` helpers + executor table; execution_matrix catalog + 3 rows; backend_manifest _ROCM_COMPILED (6+6+2 ops) + _NUMERICAL_FIXTURES; dashboards regenerated. With this, loss (29) + rl_loss (4) + quantize (8) all run on BOTH gfx1151 and AVX-512 — three of the four "0/0" categories fully cross-device. Only spectral remains. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Final ROCm-mirror PR for the loss/quant families — closes the cross-device pair for the S11 class-axis losses and the S9 quantize ops (x86 in #167/#168/#169).
class-axis loss (
rocm_class_loss_compiled, pure composition — no new C++)cross_entropy / kl / js / focal / label_smoothed_cross_entropy / z_loss; exp/log run on the rocm unary lane (gfx1151), class-axis max/sum/gather/one-hot on the host (same pattern as the x86 class-loss + ROCm attention/gemm-family lanes).fpquant + nvfp4 (new
GenerateROCMFpQuantKernel.cpp)tessera_rocm.fpquant— flat 1-operand float-grid mantissa-snap (ax=min(|x|,max);e=max(floor(log2 ax),min_exp);ulp=2^(e−mant);roundeven; sign + NaN-propagate), viamath.log2/floor/exp2/roundeven→ROCDL, parameterized for fp8 e4m3/e5m2, fp6 e2m3/e3m2, fp4 e2m1.rocm_fpquant_compiled(per-tensor scale) +rocm_nvfp4_compiled(per-block fp8-E4M3 scale + E2M1 codes, host block structure).Validation (gfx1151)
test_rocm_class_loss_compiled.py+test_rocm_fpquant_compiled.py— 15 passed vstessera.losses/tessera.ops(class-loss 2e-4; fpquant/nvfp4 2e-3 — the GPU log2/roundeven grid matches the reference on random data) + codegen gate. ruff/mypy clean; consistency green (88 passed).Wiring
runtime executors +
_rocm_fpgrid/_rocm_unary_t/_rocm_log_softmaxhelpers + executor table;execution_matrixcatalog + 3 rows;backend_manifest_ROCM_COMPILED(6+6+2 ops) +_NUMERICAL_FIXTURES; dashboards regenerated.With this, loss (29) + rl_loss (4) + quantize (8) all run on both gfx1151 and AVX-512 — three of the four "0/0" categories are now fully cross-device. Only spectral remains.
🤖 Generated with Claude Code