[aniruddhaadak80]

Summary

bit_linear and FusedBitLinear now honor the requested RMS normalization epsilon instead of silently using the fused helper's 1e-6 default. This affects direct function calls and BitNet's attention projections whenever their epsilon differs from 1e-6, including their existing 1e-8 default.

Calls follow signature order, and remaining ordinary helper calls use keyword arguments. Kernels, computation precision, public signatures, and checkpoint layout are unchanged. No test files are added or modified.

Test plan

NVIDIA H200, Python 3.12.13, PyTorch 2.9.1.post5+cu130.msh, Triton 3.5.1; flash-attn installed, FLA_CI_ENV=0, default backend dispatch enabled.

• 53 existing cases passed: 45 quantized-normalization cases, all six BitNet model cases (forward/backward, packed inputs, and generation), one BitNet fused-linear-cross-entropy case, and one BitAttention cache-guard case.
• An additional temporary numerical validation passed 72 output/gradient comparisons against a PyTorch RMSNorm + straight-through quantization reference. It covers both entry points, default/custom epsilon, FP32/FP16/BF16, optional bias, and single-row/partial-width batched inputs at the existing quantized-normalization threshold of 0.01.
• Against baseline 2c3a40f4, four default/custom-epsilon cases fail numerical comparison while both explicit-1e-6 controls pass. The fixed version passes all six.
• GPU validation used f6a909f8. Final revision d0d9bbdb removes the temporary tests and makes four equivalent positional-to-keyword substitutions; AST comparison after argument binding confirms no other executable change. Scoped pre-commit hooks and git diff --check pass; the copyright-header check also passed.

The existing affected tests can be run with:

FLA_CI_ENV=0 python -m pytest \
  tests/modules/test_layernorm.py::test_layernorm_linear_quant \
  tests/models/test_modeling_bitnet.py \
  'tests/models/test_fuse_linear_cross_entropy.py::test_fuse_linear_cross_entropy[BitNetConfig]' \
  'tests/layers/test_cu_seqlens_cache_guard.py::test_cu_seqlens_rejected_with_non_empty_cache[bitattn]' \
  -q --tb=short --maxfail=3

Other accelerator platforms and distributed execution were not tested. An initial cache-test attempt on the development host could not initialize CUDA because that host has no NVIDIA driver; the affected BitAttention case passed on H200.

Benchmark / NCU

N/A: this forwards an existing scalar argument without changing kernels, precision, tensor allocation, or launch structure. No performance improvement is claimed.

Breaking changes

None. This restores the behavior promised by the existing epsilon arguments. Outputs and gradients can change for callers that previously received the unintended 1e-6, including default BitNet projections; explicit eps=1e-6 keeps the previous behavior. No checkpoint conversion is needed.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
• Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
• This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.
• I understand this PR exceeds 500 changed lines and have explained its single purpose and why it should stay together below.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.