[adenzler-nvidia]

Sparse mass factorization launches updates at each tree depth and atomically scatters descendant contributions. This change assigns one block to each sparse tree and world, gathers each factor entry's contributions, and synchronizes between levels inside the kernel.

Model setup groups factor entries by depth and selects the block width from the mean level width. The kernel uses MuJoCo's contiguous DOF subtrees and full ancestor-chain rows. It retains the reverse LDL factor format and existing triangular solves. Dense dispatch is unchanged, sleeping trees preserve their factors, and factor storage can be strided. One block per tree limits within-tree parallelism for very large trees.

Full-step times at 8,192 worlds on an exclusive NVIDIA L40, against main at dbf69e1e:

Model	Solver	Main (ms)	This PR (ms)	Time reduction
H2 + Sharpa hands	Newton	13.138	11.138	+15.22%
H2 + Sharpa hands	CG	24.761	23.129	+6.59%
G1 flat	Newton	2.332	2.330	+0.09%
G1 flat	CG	11.733	11.686	+0.40%
Humanoid	Newton	1.260	1.260	-0.01%
Humanoid	CG	4.995	4.990	+0.12%

Medians of three alternating baseline/patched pairs of 1,000 steps, excluding setup/JIT; Warp 1.18.0, driver 615.71.10. Positive percentages mean less time. G1 and humanoid use the existing dense path and serve as controls. H2 is an assembled 81-DOF model with Sharpa hands in a zero-control contact rollout, not a policy task. G1 uses the shuffle-dance replay.

Both variants use Newton/CG iteration limits of 100/1,000 and 50 line-search iterations, with all overflow warnings enabled and treated as errors. H2 uses nconmax=128, nccdmax=128, and njmax=512 to fit the 8,192-world run in memory. All 36 rollouts completed without overflow or non-finite worlds.

Tests cover mixed dense/sparse models, wide trees, sleep/wake transitions under graph replay, deterministic mode, and strided factor storage. All pre-commit checks passed. uv run pytest -n 8 reported 3,718 passed, 42 skipped, and four failures in the existing Newton elliptic in-place Cholesky tests. All four reproduce on pristine main with the same unsupported conditional-graph memory-allocation error.

[thowell]

Thanks @adenzler-nvidia! Clean design fusing sparse $L'DL$ factorization per tree and removing per-level launches and atomic scatter updates — LGTM once two minor smooth.py items are addressed.

1. mujoco_warp/_src/smooth.py (lines 1584--1606, 1705--1719)

• Issue: Every caller of _factor_i_sparse (factor_m, factor_solve_i, and derivative.eff_metric) already dispatches _factor_blocks / _factor_solve_blocks when m.M_tiles is non-empty, which writes D for Q_LD_BLOCK_COMPACT blocks and packed Cholesky factors in L[:, :m.qLD_block_total] for dense blocks (whose entries in L[:, m.qLD_block_total:] and D are never read). Meanwhile, _factor_LD_sparse_fused reads M directly and writes all sparse-block entries of L and D. On mixed models (e.g., a sparse articulated robot plus any free object), the if m.M_tiles: branch in _factor_i_sparse adds a redundant full (nworld, nC) copy into L and a redundant (nworld, nv) _qLDiag_div launch, leaving _qLDiag_div with no other callers.
• Recommendation: Remove the if m.M_tiles: copy and _qLDiag_div launch in _factor_i_sparse, and delete the unused _qLDiag_div kernel.

def _factor_i_sparse(m: Model, d: Data, M: wp.array2d[float], L: wp.array2d[float], D: wp.array2d[float]):
  """Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
  enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
  block_dim = m.block_dim.factor_LD_sparse_fused if L.device.is_cuda else 1
  wp.launch(
    _factor_LD_sparse_fused(enable_sleep),
    dim=(d.nworld, m.qLD_fused_level_adr.size - 1, block_dim),
    inputs=[
      m.dof_treeid,
      m.M_rownnz,
      m.M_rowadr,
      m.dof_subtreenum,
      m.qLD_fused_level_adr,
      m.qLD_fused_pair_adr,
      m.qLD_fused_pairs,
      M,
      d.tree_awake,
    ],
    outputs=[L, D],
    block_dim=block_dim,
  )

2. mujoco_warp/_src/smooth.py (lines 1666--1667, 1690--1691)

• Issue: In io.py, qLD_fused_pairs groups sparse rows by depth = mjm.M_rownnz[start : start + size] - 1 == level, so M_rownnz[row] - 1 is identically equal to level (adr - level_beg, already in a register) for every pair in that level iteration. Checking entry == M_rownnz[row] - 1 issues an unnecessary global memory read in both the narrow and wide level loops.
• Recommendation: Compare entry == level directly in both branches of the level loop.

          if tid == 0:
            value = matrix[rowadr + entry] + total
            factor[rowadr + entry] = value
            if entry == level:
              diagonal[row] = 1.0 / value

          factor[rowadr + entry] = value
          if entry == level:
            diagonal[row] = 1.0 / value

[adenzler-nvidia]

Thanks @thowell, both applied in 69b8dd3.

1. Removed the if m.M_tiles: copy and the _qLDiag_div launch from _factor_i_sparse, and deleted the kernel. The solves only read D for sparse and compact dofs and the LDL region only for sparse rows, so nothing reads what those launches wrote for block dofs.
2. Both branches of the level loop now compare entry == level.