[Lucaskabela]

Purpose

Draft implementation of fused rms_norm_per_quant

Summary:

• Kernel: scale_ue8m0 template flag on compute_dynamic_per_token_scales (scalar + vectorized).
  Power-of-two rounding extracted into adjust_scale_for_ue8m0 device inline helper with a named
  MIN_SCALE_FOR_LOG = 1e-10f constant; both call sites use the helper.
• Op schema: rms_norm_per_block_quant gains bool scale_ue8m0 (csrc/ops.h, torch_bindings.cpp,
  _custom_ops.py).
• Routing: rms_quant_fusion.py propagates scale_ue8m0=self.is_e8m0; input_quant_fp8.py adds
  _skip_packed_for_fusion so fusion can match when fuse_norm_quant is on.

Test Plan

pytest tests/compile/passes/test_fusion.py::test_fusion_rmsnorm_quant

We also run a benchmark of this naive impl (fusion) vs the packed version

Test Result

That is fair - here is the timing result:

N	D	packed unfused (µs)	fused UE8M0 (µs)	winner	Δ
1	4096	38.1	30.0	fused	−21%
1	7168	38.1	30.0	fused	−21%
128	4096	39.2	30.5	fused	−22%
128	7168	39.2	31.0	fused	−21%
1024	4096	38.5	30.5	fused	−21%
1024	7168	38.9	34.1	fused	−12%
2048	4096	38.8	34.8	fused	−10%
2048	7168	46.9	55.4	packed	+18%
3072	4096	42.6	49.3	packed	+16%
3072	7168	71.8	84.3	packed	+17%
4096	4096	53.4	65.7	packed	+23%
4096	7168	98.7	112.4	packed	+14%
8192	4096	114.5	130.4	packed	+14%
8192	7168	195.7	218.7	packed	+12%

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request introduces support for UE8M0 scaling in the rms_norm_per_block_quant kernel and its associated fusion passes. It adds a scale_ue8m0 flag across the C++ kernels, PyTorch bindings, and Python wrappers, implementing power-of-two rounding for scales when enabled. Additionally, the InputQuantFP8 layer is updated to coordinate between packed UE8M0 scales and the norm-quant fusion pass. Feedback was provided regarding the duplication of the scale adjustment logic in layernorm_utils.cuh and the absence of a centralized helper function or constant for the rounding calculation.