Remove the view() calls from Linear/LNLinear/LNMLP modules

[ptrendx]

Description

This PR removes the unnecessary calls to view() from the TE modules to avoid the CPU overhead.

Results on GB200 (just the CPU time is reported):

Module             Precision               Forward    CPU reduction    Forward + backward    CPU reduction
  ━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━
   Linear             BF16           59.19 → 51.19 µs            13.5%    311.30 → 276.50 µs            11.2%
  ─────────────────  ───────────  ────────────────────  ───────────────  ────────────────────  ───────────────
   Linear             MXFP8        114.51 → 103.62 µs             9.5%    440.93 → 416.92 µs             5.4%
  ─────────────────  ───────────  ────────────────────  ───────────────  ────────────────────  ───────────────
   LayerNormLinear    BF16           83.48 → 65.25 µs            21.8%    397.73 → 349.00 µs            12.3%
  ─────────────────  ───────────  ────────────────────  ───────────────  ────────────────────  ───────────────
   LayerNormLinear    MXFP8        135.42 → 121.07 µs            10.6%    518.67 → 490.49 µs             5.4%
  ─────────────────  ───────────  ────────────────────  ───────────────  ────────────────────  ───────────────
   LayerNormMLP       BF16         210.44 → 189.08 µs            10.2%    749.00 → 683.15 µs             8.8%
  ─────────────────  ───────────  ────────────────────  ───────────────  ────────────────────  ───────────────
   LayerNormMLP       MXFP8        236.67 → 221.99 µs             6.2%    797.14 → 765.84 µs             3.9%

Each commit is standalone and can be reviewed on its own.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ptrendx]

/te-ci

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding findings or newly introduced actionable failures.

Summary

This PR removes activation-flattening and output-restoration views from Linear, LayerNormLinear, and LayerNormMLP while extending lower-level normalization, transpose, GEMM, bias-gradient, custom-recipe, and communication-overlap paths to preserve logical tensor shapes directly.

• Preserves rank-1 and higher-rank activation shapes through forward and backward operations.
• Updates normalization, transpose, GEMM, and bias-gradient implementations for logical N-D tensors.
• Adds coverage for noncontiguous inputs, custom recipes, distributed overlap, and torch.compile.
• The only change since the previous review adds ndim support to debug quantized tensors.

_{Reviews (4) · Last reviewed commit: "Fix the debug test"}

[ptrendx]

/te-ci

[ptrendx]

/te-ci

[vthumbe1503]

LGTM. My comments are nits. Its surprising how using views for semantic convinience introduces so much CPU overheads. GroupedLinear would also benefit with similar optimizations I hope.

[vthumbe1503]

Nit: can we use flat_2d_dims for output as well?

Suggested change

	NVTE_CHECK(output.data.shape[0] == row_length, "Output has invalid first dimension (expected ",
	row_length, ", got ", output.data.shape[0], ").");
	const size_t output_cols = std::accumulate(output.data.shape.begin() + 1, output.data.shape.end(),
	size_t{1}, std::multiplies<size_t>());
	NVTE_CHECK(output_cols == num_rows,
	"Output has invalid product of dimensions after the first (expected ", num_rows,
	", got ", output_cols, ").");
	const auto [out_rows, out_cols] = output.flat_2d_dims();
	NVTE_CHECK(out_rows== row_length, "Output has invalid first dimension (expected ",
	row_length, ", got ", out_rows, ").");
	NVTE_CHECK(out_cols == num_rows,
	"Output has invalid product of dimensions after the first (expected ", num_rows,
	", got ", out_cols, ").");

[vthumbe1503]

lint seems to be failing since tp_size is not longer used in this function