[PyTorch] [torch.compile] torch.compile support for GroupedLinear

[pggPL]

Description

Follow-up to NVIDIA#3053 (linear_torch_compile_final_attempt, torch.compile custom-op support for Linear). This PR adds the same torch.compile custom-op support to te.pytorch.GroupedLinear: torch.compile(fullgraph=True) traces the module as a single custom-op graph node under every built-in recipe, in default and reduce-overhead (CUDA-graph trees) modes.

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

Framework extensions in dynamo/custom_op.py

GroupedLinear passes num_gemms weights / biases / cached FP8 workspaces and 6 quantizer lists; the custom-op framework only supported scalar fields. Added (all additive; the Linear op registration is untouched):

• _ListAdapter — a generic lift of any scalar field adapter to its List[...] analog. The schema is fixed at registration while list lengths vary per call, so the inner adapter's slots are aggregated instead of replicated: Tensor/Tensor? → Tensor[] (one entry per element, None rides the existing 0-element complex32 sentinel), Tensor[] → Tensor[] (elements' payloads concatenated in order, per-element counts in the bundle slot), value bundle → one bundle of per-element bundles. to_slots/from_slots delegate per element to the scalar adapter, so kind tagging / storage flattening / value-opaque checks apply unchanged. Covers List[Tensor], List[Tensor | QuantizedTensorStorage | None] (weights, workspaces) and List[Quantizer]; any future scalar adapter becomes list-capable for free.
• _SymIntListAdapter — List[int] as a SymInt[] slot: host-side ints that Dynamo may mark dynamic (e.g. m_splits) cross as op inputs instead of being baked into a value bundle, so one symbolic graph serves every split distribution.
• List-aware grad plumbing — _resolve_grad_targets returns (slot, is_list) targets; _pack_bwd_result flattens a list-shaped grad to one slot per element; _autograd_backward slices the flat bwd return back per target (with a count check); fwd_tensor_list_lengths now records only genuine tensor lists (a SymInt[] slot takes a plain None grad).
• Subclass flattening for list slots — the wrapper op's register_torch_dispatch path rewrites Float8Tensor-style entries inside a list group into the storage layout (splicing inner tensors into the concatenated slot and rebuilding the meta bundle).

grouped_linear.py refactor + op registration

Mirrors the linear.py structure:

• GroupedLinearFwdArgs / GroupedLinearBwdArgs dataclasses; the save_original_input runtime flips (delayed-scaling / unsafe-requantization warnings) and backward_override are resolved in the module forward, before the args are built, so the impl and the compile-time fake always see one final flag. Recipe-derived scalars (fprop/dgrad/wgrad split accumulators, native-bgrad support) are computed there too — the backward no longer reads the recipe object (_split_quantize_and_bias takes a precomputed bool instead of recipe).
• _grouped_linear_forward_impl / _grouped_linear_backward_impl extracted from the autograd.Function staticmethods; _GroupedLinear is now a thin wrapper (the fused GroupedTensor path dispatches before the impl and is unchanged). _grouped_linear_setup_ctx populates the backward args and rebuilds alias-deduped save slots; _grouped_linear_eager (@no_torch_dynamo) is the eager entry.
• Saved-tensor layout (inputmat_full, *inputmats, *weights_fp8, *saved_weights, *biases) (1 + 4N slots) with per-slot name/index alias tags — a custom op may not return aliasing tensors, and every weights_fp8[i] aliases an input or another return (cache hit → weight_workspaces[i], miss → new_workspaces[i], primary FP8 → weights[i]). The plain (non-quantized) path saves one full cast buffer (or aliases inp for a no-op cast) and backward re-splits it — the N torch.split views can't be returned from an op.
• On the compiled path only (compiled_op=True): split-quantize bulk allocation is disabled (bulk storages alias each other) and backward allocates per-expert wgrad buffers instead of one packed buffer, for the same reason; clear_tensor_data on saved inputs is skipped (in-place storage resize of op inputs).
• _grouped_linear_forward_fake / _grouped_linear_backward_fake — TensorSpec twins mirroring the impl's quantizer set_usage / update_usage retargeting and the alias tags exactly.
• compile_unsupported_reason() gate (explicit torch._dynamo.graph_break(msg=...) then warn_compile_eager_fallback, as in Linear): debug, distributed weights, m_splits passed as a device tensor, user out/dgrad_out buffers (returned alias), fused GroupedTensor path, single_grouped_weight/bias, CPU offloading, delayed wgrad, fuse_wgrad_accumulation, fp8_calibration (calibration inside the op would mutate quantizers rebuilt from graph constants), mixed requires_grad across the weights list, non-value-opaque quantizers.
• GroupedLinear.forward drops @no_torch_dynamo(); the compiled path requires host-side m_splits (list/tuple — a device tensor falls back), runs a grouped check_gemm_dims variant (check_grouped_gemm_dims in utils.py, torch._check guards incl. per-split FP8 divisibility), and reassembles new_workspaces from the op outputs to refresh _fp8_workspaces.
• GroupedLinear.reset_parameters eagerly pre-allocates the grouped cuBLAS workspace (get_cublas_workspace(dev, False, True)), mirroring Linear.reset_parameters — under reduce-overhead the first grouped GEMM can run inside CUDA-graph capture.

Deviations from the original plan

• Dynamic m_splits are supported, not deferred: instead of baking the splits into the OpaqueValueBundle (graph-per-distribution, and a hard Dynamo error once automatic-dynamic marks the ints symbolic), m_splits crosses as a SymInt[] op input. Static splits behave as before; once a distribution changes, Dynamo recompiles once to a symbolic graph that then serves all further distributions.
• fp8_calibration resolved to a gate entry (open question 2): calibrate() inside the op body would update amax on quantizers rebuilt from __fx_repr__ constants and the updates would be lost.
• Uniform requires_grad across the weights list is asserted in the gate (open question 3).

Tests (tests/pytorch/test_torch_compile.py)

• test_te_grouped_linear_compiles — fullgraph, all recipes + bf16 × {default, reduce-overhead}, bit-exact eager-vs-compiled forward / dgrad / per-expert wgrads + bias grads, cudagraph-skip counter guard.
• test_te_grouped_linear_compile_with_quantized_fp8_weight — primary-FP8 weights (subclass flattening inside a Tensor[] slot group).
• test_te_grouped_linear_compile_is_first_microbatch — per-expert FP8 workspace caching across a [True, False, False] schedule with cache-identity checks, vs a separate eager reference module.
• test_te_grouped_linear_m_splits_change — same splits: no recompile; changed distributions: still the compiled path (symbolic graph), bit-exact.
• test_te_grouped_linear_compile_save_original_input — input saved by alias, re-split/re-quantized in backward.
• test_te_grouped_linear_compile_train_eval_switch — grouped counterpart of the Linear xfail (currently xpasses).
• test_te_grouped_linear_compile_m_splits_tensor_falls_back / ..._delayed_scaling_falls_back — graph-break fallbacks; with fullgraph=True the error names the reason.
• test_custom_op_tensor_list_grads — framework unit test: a toy op with a List[Optional[Tensor]] input (incl. a None entry riding the sentinel) compiles fullgraph and routes one gradient per list element, eager and compiled.

Test results (RTX Ada workstation)

• test_torch_compile.py: 69 passed, 24 skipped, 1 xfailed, 1 xpassed
• test_grouped_linear.py (eager): 717 passed, 416 skipped — the refactor is a numerical no-op for eager

Known limitations / follow-ups

1. Fused GroupedTensor path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1) stays eager (gate); it keeps m_splits on device and is the natural follow-up now that SymInt[] splits work.
2. Compiled-path split-quantize runs with bulk allocation disabled (returned storages must not alias); if the extra allocs show up in profiles, a packed-buffer save + in-op re-split is the follow-up.
3. fuse_wgrad_accumulation, delayed wgrad (wgrad_store), CPU offloading, debug and out=/dgrad_out= buffers fall back to eager, same as (or analogous to) Linear.
4. Quantizer constants scale with 6×N per graph; fine at tested expert counts, worth a compile-time measurement at large N.

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