Add dist GEMM attention and FFN backend by anijain2305 · Pull Request #4044 · pytorch/torchtitan

anijain2305 · 2026-07-31T23:26:14Z

Stack from ghstack (oldest at bottom):

Motivation

Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM
rather than running them beside it -- is currently only reachable in torchtitan
through a compiler: either torch.compile's micro-pipelining pass behind
--parallelism.enable_async_tensor_parallel, or a GraphTrainer graph pass. That
ties a parallelism strategy to a particular compilation path.

This makes it a property of the model config instead, selected by a
gemm_backend argument on model_registry exactly as attn_backend already
selects the attention kernel. Async-TP then works under the eager trainer with no
graph pass involved.

Structure

torchtitan/distributed/linear.py holds three reusable autograd Functions.
AllGatherLinear is a column-parallel projection (all-gather the sequence shard,
then GEMM) and LinearReduceScatter its row-parallel dual (GEMM, then
reduce-scatter). AllGatherLinearMulti serves several column-parallel linears
that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than
two; its backward still needs only the same two collectives, because the sum over
outputs is expressed as a single concatenated product. None of them is
attention-specific; MoE projections could use the same pair.

torchtitan/models/common/dist_gemm.py holds the modules.
AllGatherFusedQKVLinear, RowParallelLinear and AllGatherFusedFeedForward are
drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the
stock parameter layouts so checkpoints still interoperate. RowParallelLinear
serves both attention's wo and the FFN's w2.

Selection follows the attn_backend path one for one: model_registry ->
flavor function -> _build_llama3_layers -> make_gqa_config / make_ffn_config,
which is where gemm_backend="dist_gemm" picks the fused classes. Only llama3 is
threaded; every other model keeps the default and can gain the argument when
someone needs it. llama3_debugmodel_dist_gemm is the config-registry entry, and
the h100 integration test runs it at TP=2.

No attention or FFN subclass is needed beyond the projections themselves. That
required one fix in shared code: GQAttention.forward captured B, L from its
input and reused them to fold the attention output, which is only valid when the
QKV preserves sequence length. AllGatherFusedQKVLinear gathers the SP shard, so
it does not, and under spmd_types -- where a tensor's shape is its local shape --
that silently folded sequence into features. It now derives the fold from the
tensor being folded.

spmd_types only

The fused modules take and return plain local tensors, which is the spmd_types
contract; the DTensor backends are being deprecated and are rejected outright.
Two preconditions are checked at parallelize time, because neither is detectable
from inside a module once activations carry no placements: the backend must be
spmd_types, and sequence parallelism must be enabled. The fused GEMMs are the
SP collectives, so with SP disabled there is nothing to gather and wo would
reduce-scatter where it must all-reduce.

Because selection happens at config-construction time, set_gqa_attention_sharding
and set_dense_ffn_sharding declare the fused blocks' contracts directly: no
boundary all-gather, and a row-parallel output that emits its final Shard(1)
rather than the Partial a stock rowwise linear produces. Both branches are
transitional and collapse once redistribute collectives move inside the modules
generally.

Symmetric memory

One workspace per process group serves every layer, and the symm_mem ops size it
themselves: each computes its own requirement and get_symm_mem_workspace grows
monotonically, so the cost is a max over layers rather than a sum. Nothing here
reserves it up front. Growth re-rendezvouses, which is a collective and is
rejected during CUDA graph capture, but capture is always preceded by a real
warmup call (graph_trainer/cudagraph.py), and a warmup forward and backward runs
every layer's collectives -- so all growth has happened before anything is
recorded. An earlier revision pre-reserved from parallelize to make that
explicit; it bought one growth instead of a few during warmup, at the cost of
duplicating the sizing arithmetic, and is not what makes capture safe.

One consequence worth keeping: both wgrad all-gathers leave return_A at its
default. Passing False would select _multimem_all_gather_matmul, which reserves
the full gathered buffer rather than a shard -- ranks times more symmetric memory
for that one call -- and its heuristic evaluates to K <= 2048 there, a dimension
it was not tuned for. Leaving it defaulted also keeps forward and backward on the
same schedule.

Every op carves its buffers from offset 0 of that shared workspace, so the regions
alias; safety comes from the ops being serialized rather than from disjointness.
Each op also brackets itself with barriers, so ranks cannot run ahead of one
another. Sequential module forwards and autograd backward are single-stream, which
is how models actually run, so this holds today -- but deliberate cross-stream
overlap would need distinct workspace offsets, not a barrier.

Known gap

--debug.spmd_typechecking does not pass. The autograd Functions are registered
with typecheck_forward declaring their type transitions, which the checker
requires, but it then rejects the QKV reshape that splits the TP-sharded feature
dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over
n_kv. The stock FusedQKVLinear hits the same limitation and works around it with
spmd.local() plus an explicit PartitionSpec, carrying a TODO to remove that
once the checker handles aligned splits. Rather than add a second instance of a
pattern already marked for deletion, this leaves typechecking unsupported until
that lands upstream.

Authored with assistance from Claude Code.

[ghstack-poisoned]

ghstack-source-id: debdd0f Pull-Request: #4044

[ghstack-poisoned]

ghstack-source-id: aee34d3 Pull-Request: #4044

[ghstack-poisoned]

Wire the symmetric-memory fused all-gather-matmul and matmul-reduce-scatter ops into GQA attention, so the TP collective around each projection is folded into the GEMM instead of running beside it. Two reusable autograd Functions in `torchtitan/distributed/dist_linear.py` do the work: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter). Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring: `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate, and `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Two things in here are less obvious than the rest and are worth reviewing closely. Row ordering. The all-gather concatenates whole per-rank blocks, so the rows it produces are ordered (rank, batch, seq_local). Flattening [B, S/W, D] directly would be reinterpreted as (batch, seq) and would mix batches together once bsz > 1, so both forwards put the sequence outermost before flattening. The reduce-scatter side has the mirror problem: it splits the flattened rows, so the sequence has to be outermost or the split cuts across batches instead of the sequence. Feeding 2D with scatter_dim=0 is also what lets the operator take its fused schedules, which it declines for a 3D input. The wo output contract. A stock rowwise linear emits a Partial sum over its slice of K and lets the framework reduce-scatter it, which is what `set_gqa_attention_sharding` declares via `rowwise_config()`. `AttentionOutputLinear` collapses those two steps -- the reduce-scatter happens inside the fused op -- so its forward returns Shard(1) directly and never produces a Partial. Left alone the module fails its own out_src check. The override factory does install a corrected config, but the model's sharding setup runs afterwards and overwrites it, so the correction is applied in `DistGemmGQAttention.parallelize` instead, after the point where it would have been clobbered. Symmetric-memory workspace. The fused ops grow their workspace lazily, and growing it re-rendezvouses, which is a collective and is rejected during CUDA graph capture; worse, growth frees the old buffer while its address stays baked into any graph already captured against it. `maybe_update_dist_gemm_config` therefore stamps the per-step token count from the runtime config onto the module configs, and `parallelize` reserves the worst case before any layer runs. This does not remove the need for a warmup step -- the reservation happens inside `forward` -- but it means one warmup step reaches the final size instead of creeping up as different paths first execute. Test plan: Unit tests, 8 passing on 2 GPUs: ``` python -m pytest tests/unit_tests/test_dist_linear.py tests/unit_tests/test_dist_gemm_attention.py -q ``` ``` 8 passed ``` `test_dist_linear.py` checks the two Functions against a single-device reference built from the unsharded weights: forward, dgrad, wgrad and bias. The terms that involve no cross-rank reduction are asserted bit-exact and only the reduced ones get a bf16 tolerance, so a transposed or mis-sharded gradient cannot hide inside a loose comparison. `test_dist_gemm_attention.py` has four CPU config-graph tests (does the override rewrite the graph, does the runtime config reach it, is a non-Trainer config a no-op) and one 2-GPU test pinning the wo contract fix described above. Note the replacement modules cannot be tested standalone: both the weight sharding and the output contract are written onto their configs by the model's sharding setup, so a bare `AttentionOutputLinear.Config(...).build()` has an unsharded weight and no contract at all. The GPU test goes through the attention block for that reason, and says so, since it is the first thing anyone extending this will trip over. Integration test, added to the h100 suite: ``` python -m tests.integration_tests.run_tests --test_suite h100 \ --test_name override_dist_gemm_attention --ngpu 4 <out> ``` ``` step: 10 loss: 4.25631 grad_norm: 2.1551 ... rc=0 ``` It lives in `h100.py` next to the existing async-TP entry rather than in `features.py`, on the same hardware class, and because `h100.py` returns its entries unwrapped -- `features.py` would auto-generate a `--debug.spmd_typechecking` variant, which this override would fail by construction, since the fix above deliberately removes a declared placement. Confirmed the fused paths actually execute rather than silently falling back to `super().forward()`: wrapping the two Functions and running one fwd+bwd of the 6-layer llama3 debugmodel at TP=2 counts 6 `AllGatherLinear` and 6 `LinearReduceScatter` invocations, one of each per layer. Authored with assistance from Claude Code. ghstack-source-id: 184161d Pull-Request: #4044

[ghstack-poisoned]

Wire the symmetric-memory fused all-gather-matmul and matmul-reduce-scatter ops into GQA attention, so the TP collective around each projection is folded into the GEMM instead of running beside it. Two reusable autograd Functions in `torchtitan/distributed/dist_linear.py` do the work: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter). Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring: `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate, and `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Two things in here are less obvious than the rest and are worth reviewing closely. Row ordering. The all-gather concatenates whole per-rank blocks, so the rows it produces are ordered (rank, batch, seq_local). Flattening [B, S/W, D] directly would be reinterpreted as (batch, seq) and would mix batches together once bsz > 1, so both forwards put the sequence outermost before flattening. The reduce-scatter side has the mirror problem: it splits the flattened rows, so the sequence has to be outermost or the split cuts across batches instead of the sequence. Feeding 2D with scatter_dim=0 is also what lets the operator take its fused schedules, which it declines for a 3D input. The wo output contract. A stock rowwise linear emits a Partial sum over its slice of K and lets the framework reduce-scatter it, which is what `set_gqa_attention_sharding` declares via `rowwise_config()`. `AttentionOutputLinear` collapses those two steps -- the reduce-scatter happens inside the fused op -- so its forward returns Shard(1) directly and never produces a Partial. Left alone the module fails its own out_src check. The override factory does install a corrected config, but the model's sharding setup runs afterwards and overwrites it, so the correction is applied in `DistGemmGQAttention.parallelize` instead, after the point where it would have been clobbered. Symmetric-memory workspace. The fused ops grow their workspace lazily, and growing it re-rendezvouses, which is a collective and is rejected during CUDA graph capture; worse, growth frees the old buffer while its address stays baked into any graph already captured against it. `maybe_update_dist_gemm_config` therefore stamps the per-step token count from the runtime config onto the module configs, and `parallelize` reserves the worst case before any layer runs. This does not remove the need for a warmup step -- the reservation happens inside `forward` -- but it means one warmup step reaches the final size instead of creeping up as different paths first execute. Test plan: Unit tests, 8 passing on 2 GPUs: ``` python -m pytest tests/unit_tests/test_dist_linear.py tests/unit_tests/test_dist_gemm_attention.py -q ``` ``` 8 passed ``` `test_dist_linear.py` checks the two Functions against a single-device reference built from the unsharded weights: forward, dgrad, wgrad and bias. The terms that involve no cross-rank reduction are asserted bit-exact and only the reduced ones get a bf16 tolerance, so a transposed or mis-sharded gradient cannot hide inside a loose comparison. `test_dist_gemm_attention.py` has four CPU config-graph tests (does the override rewrite the graph, does the runtime config reach it, is a non-Trainer config a no-op) and one 2-GPU test pinning the wo contract fix described above. Note the replacement modules cannot be tested standalone: both the weight sharding and the output contract are written onto their configs by the model's sharding setup, so a bare `AttentionOutputLinear.Config(...).build()` has an unsharded weight and no contract at all. The GPU test goes through the attention block for that reason, and says so, since it is the first thing anyone extending this will trip over. Integration test, added to the h100 suite: ``` python -m tests.integration_tests.run_tests --test_suite h100 \ --test_name override_dist_gemm_attention --ngpu 4 <out> ``` ``` step: 10 loss: 4.25631 grad_norm: 2.1551 ... rc=0 ``` It lives in `h100.py` next to the existing async-TP entry rather than in `features.py`, on the same hardware class, and because `h100.py` returns its entries unwrapped -- `features.py` would auto-generate a `--debug.spmd_typechecking` variant, which this override would fail by construction, since the fix above deliberately removes a declared placement. Confirmed the fused paths actually execute rather than silently falling back to `super().forward()`: wrapping the two Functions and running one fwd+bwd of the 6-layer llama3 debugmodel at TP=2 counts 6 `AllGatherLinear` and 6 `LinearReduceScatter` invocations, one of each per layer. Authored with assistance from Claude Code. ghstack-source-id: 158e897 Pull-Request: #4044

[ghstack-poisoned]

Wire the symmetric-memory fused all-gather-matmul and matmul-reduce-scatter ops into GQA attention, so the TP collective around each projection is folded into the GEMM instead of running beside it. Two reusable autograd Functions in `torchtitan/distributed/dist_linear.py` do the work: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter). Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring: `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate, and `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Two things in here are less obvious than the rest and are worth reviewing closely. Row ordering. The all-gather concatenates whole per-rank blocks, so the rows it produces are ordered (rank, batch, seq_local). Flattening [B, S/W, D] directly would be reinterpreted as (batch, seq) and would mix batches together once bsz > 1, so both forwards put the sequence outermost before flattening. The reduce-scatter side has the mirror problem: it splits the flattened rows, so the sequence has to be outermost or the split cuts across batches instead of the sequence. Feeding 2D with scatter_dim=0 is also what lets the operator take its fused schedules, which it declines for a 3D input. The wo output contract. A stock rowwise linear emits a Partial sum over its slice of K and lets the framework reduce-scatter it, which is what `set_gqa_attention_sharding` declares via `rowwise_config()`. `AttentionOutputLinear` collapses those two steps -- the reduce-scatter happens inside the fused op -- so its forward returns Shard(1) directly and never produces a Partial. Left alone the module fails its own out_src check. The override factory does install a corrected config, but the model's sharding setup runs afterwards and overwrites it, so the correction is applied in `DistGemmGQAttention.parallelize` instead, after the point where it would have been clobbered. Symmetric-memory workspace. The fused ops grow their workspace lazily, and growing it re-rendezvouses, which is a collective and is rejected during CUDA graph capture; worse, growth frees the old buffer while its address stays baked into any graph already captured against it. `maybe_update_dist_gemm_config` therefore stamps the per-step token count from the runtime config onto the module configs, and `parallelize` reserves the worst case before any layer runs. This does not remove the need for a warmup step -- the reservation happens inside `forward` -- but it means one warmup step reaches the final size instead of creeping up as different paths first execute. Test plan: Unit tests, 8 passing on 2 GPUs: ``` python -m pytest tests/unit_tests/test_dist_linear.py tests/unit_tests/test_dist_gemm_attention.py -q ``` ``` 8 passed ``` `test_dist_linear.py` checks the two Functions against a single-device reference built from the unsharded weights: forward, dgrad, wgrad and bias. The terms that involve no cross-rank reduction are asserted bit-exact and only the reduced ones get a bf16 tolerance, so a transposed or mis-sharded gradient cannot hide inside a loose comparison. `test_dist_gemm_attention.py` has four CPU config-graph tests (does the override rewrite the graph, does the runtime config reach it, is a non-Trainer config a no-op) and one 2-GPU test pinning the wo contract fix described above. Note the replacement modules cannot be tested standalone: both the weight sharding and the output contract are written onto their configs by the model's sharding setup, so a bare `AttentionOutputLinear.Config(...).build()` has an unsharded weight and no contract at all. The GPU test goes through the attention block for that reason, and says so, since it is the first thing anyone extending this will trip over. Integration test, added to the h100 suite: ``` python -m tests.integration_tests.run_tests --test_suite h100 \ --test_name override_dist_gemm_attention --ngpu 4 <out> ``` ``` step: 10 loss: 4.25631 grad_norm: 2.1551 ... rc=0 ``` It lives in `h100.py` next to the existing async-TP entry rather than in `features.py`, on the same hardware class, and because `h100.py` returns its entries unwrapped -- `features.py` would auto-generate a `--debug.spmd_typechecking` variant, which this override would fail by construction, since the fix above deliberately removes a declared placement. Confirmed the fused paths actually execute rather than silently falling back to `super().forward()`: wrapping the two Functions and running one fwd+bwd of the 6-layer llama3 debugmodel at TP=2 counts 6 `AllGatherLinear` and 6 `LinearReduceScatter` invocations, one of each per layer. Authored with assistance from Claude Code. ghstack-source-id: a395aa5 Pull-Request: #4044

[ghstack-poisoned]

Wire the symmetric-memory fused all-gather-matmul and matmul-reduce-scatter ops into GQA attention, so the TP collective around each projection is folded into the GEMM instead of running beside it. Two reusable autograd Functions in `torchtitan/distributed/dist_linear.py` do the work: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter). Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring: `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate, and `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Two things in here are less obvious than the rest and are worth reviewing closely. Row ordering. The all-gather concatenates whole per-rank blocks, so the rows it produces are ordered (rank, batch, seq_local). Flattening [B, S/W, D] directly would be reinterpreted as (batch, seq) and would mix batches together once bsz > 1, so both forwards put the sequence outermost before flattening. The reduce-scatter side has the mirror problem: it splits the flattened rows, so the sequence has to be outermost or the split cuts across batches instead of the sequence. Feeding 2D with scatter_dim=0 is also what lets the operator take its fused schedules, which it declines for a 3D input. The wo output contract. A stock rowwise linear emits a Partial sum over its slice of K and lets the framework reduce-scatter it, which is what `set_gqa_attention_sharding` declares via `rowwise_config()`. `AttentionOutputLinear` collapses those two steps -- the reduce-scatter happens inside the fused op -- so its forward returns Shard(1) directly and never produces a Partial. Left alone the module fails its own out_src check. The override factory does install a corrected config, but the model's sharding setup runs afterwards and overwrites it, so the correction is applied in `DistGemmGQAttention.parallelize` instead, after the point where it would have been clobbered. Symmetric-memory workspace. The fused ops grow their workspace lazily, and growing it re-rendezvouses, which is a collective and is rejected during CUDA graph capture; worse, growth frees the old buffer while its address stays baked into any graph already captured against it. `maybe_update_dist_gemm_config` therefore stamps the per-step token count from the runtime config onto the module configs, and `parallelize` reserves the worst case before any layer runs. This does not remove the need for a warmup step -- the reservation happens inside `forward` -- but it means one warmup step reaches the final size instead of creeping up as different paths first execute. Test plan: Unit tests, 8 passing on 2 GPUs: ``` python -m pytest tests/unit_tests/test_dist_linear.py tests/unit_tests/test_dist_gemm_attention.py -q ``` ``` 8 passed ``` `test_dist_linear.py` checks the two Functions against a single-device reference built from the unsharded weights: forward, dgrad, wgrad and bias. The terms that involve no cross-rank reduction are asserted bit-exact and only the reduced ones get a bf16 tolerance, so a transposed or mis-sharded gradient cannot hide inside a loose comparison. `test_dist_gemm_attention.py` has four CPU config-graph tests (does the override rewrite the graph, does the runtime config reach it, is a non-Trainer config a no-op) and one 2-GPU test pinning the wo contract fix described above. Note the replacement modules cannot be tested standalone: both the weight sharding and the output contract are written onto their configs by the model's sharding setup, so a bare `AttentionOutputLinear.Config(...).build()` has an unsharded weight and no contract at all. The GPU test goes through the attention block for that reason, and says so, since it is the first thing anyone extending this will trip over. Integration test, added to the h100 suite: ``` python -m tests.integration_tests.run_tests --test_suite h100 \ --test_name override_dist_gemm_attention --ngpu 4 <out> ``` ``` step: 10 loss: 4.25631 grad_norm: 2.1551 ... rc=0 ``` It lives in `h100.py` next to the existing async-TP entry rather than in `features.py`, on the same hardware class, and because `h100.py` returns its entries unwrapped -- `features.py` would auto-generate a `--debug.spmd_typechecking` variant, which this override would fail by construction, since the fix above deliberately removes a declared placement. Confirmed the fused paths actually execute rather than silently falling back to `super().forward()`: wrapping the two Functions and running one fwd+bwd of the 6-layer llama3 debugmodel at TP=2 counts 6 `AllGatherLinear` and 6 `LinearReduceScatter` invocations, one of each per layer. Authored with assistance from Claude Code. ghstack-source-id: 09b54c2 Pull-Request: #4044

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the module instead. The fused projections are ordinary `Module`s selected through `override.imports`, so async-TP works under the eager trainer with no graph pass involved, composes with FSDP and TP through the usual `ShardingConfig`, and requires no edits to any model's config registry. It also serves as the first override example that fuses a collective, alongside the existing parametrization (`fused_swiglu`) and custom-kernel (`helion_rope`) ones. Structure `torchtitan/distributed/dist_linear.py` holds two reusable autograd Functions: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter), plus `reserve_symm_mem_workspace`. Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring. `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate. `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Symmetric memory There is one workspace per process group, shared by every layer and grown to the max over layers rather than their sum, so the cost is the largest layer once. It is reserved from `parallelize`, before any layer runs. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address stays baked into any graph already captured against it -- a use-after-free on replay rather than a clean error. Reserving every layer up front makes that unreachable. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Authored with assistance from Claude Code. ghstack-source-id: 597a58a Pull-Request: #4044

tianyu-l · 2026-08-11T21:11:52Z

let's not use override -- override is mainly for out-of-repo development / short-term unblocking work

DistGEMM is something we want to maintain, so let's support properly.

We can have a Literal arg like gemm_backend in model_registry to select from a set of options including dist_gemm

Refactored as per the suggestion.

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the module instead. The fused projections are ordinary `Module`s selected through `override.imports`, so async-TP works under the eager trainer with no graph pass involved, composes with FSDP and TP through the usual `ShardingConfig`, and requires no edits to any model's config registry. It also serves as the first override example that fuses a collective, alongside the existing parametrization (`fused_swiglu`) and custom-kernel (`helion_rope`) ones. Structure `torchtitan/distributed/dist_linear.py` holds two reusable autograd Functions: `AllGatherLinear` for a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` for a row-parallel one (GEMM, then reduce-scatter), plus `reserve_symm_mem_workspace`. Neither is attention-specific -- FFN and MoE projections could use the same pair -- so they live under `torchtitan/distributed` rather than beside the override. `torchtitan/overrides/dist_gemm_attention.py` holds the wiring. `AllGatherFusedQKVLinear` and `AttentionOutputLinear` are drop-in replacements for the stock QKV and output projections, keeping the stock parameter layouts so checkpoints still interoperate. `DistGemmGQAttention` wires both into `GQAttention` and drops the parent attention-boundary all-gather that the fused QKV now owns. Symmetric memory There is one workspace per process group, shared by every layer and grown to the max over layers rather than their sum, so the cost is the largest layer once. It is reserved from `parallelize`, before any layer runs. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address stays baked into any graph already captured against it -- a use-after-free on replay rather than a clean error. Reserving every layer up front makes that unreachable. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Authored with assistance from Claude Code. ghstack-source-id: 388be90 Pull-Request: #4044

[ghstack-poisoned]

sanketpurandare · 2026-08-12T00:45:32Z

+        # return_A=False matters beyond saving the copy: asking for the gathered
+        # tensor back disqualifies the op's multimem fast path, which is the only
+        # path that beats an unfused all-gather + mm at small token counts.
+        _, grad_w_outputs = torch.ops.symm_mem.fused_all_gather_matmul(


AllGatherLinear.backward() deliberately passes return_A=False, which allows
PyTorch to choose _multimem_all_gather_matmul on multicast-capable devices.
That schedule reserves the full gathered A buffer. For the wgrad call, its shape
is [K, M], not the local [K / W, M] shard assumed by this reservation.

For a layer where K is max(in_features, out_features), the pre-reservation
is:

2 * (M / W) * K * sizeof(float32) = 8 * M * K / W bytes

The selected multimem schedule requests M * K * input.element_size() bytes.
The reservation is therefore only half the BF16 requirement at TP=8; for FP32
it is insufficient above TP=2. During CUDA graph capture, PyTorch rejects the
resulting workspace growth. Outside capture, the lazy growth also invalidates
the stated guarantee that the workspace has reached its final size before any
layer runs. Llama 3 1B is a concrete affected shape: K=2048, its TP=8 QKV
output shard is smaller than K, and PyTorch's global K <= 2048 heuristic
selects the multimem path on a multicast-capable device.

Fixed it by setting return_A to True.

sanketpurandare · 2026-08-12T00:48:49Z

Some agent generated stuff:

1. Resolve or characterize the numerical drift

A deterministic TP=2 comparison of stock attention versus dist_gemm produced
bitwise-identical loss for two steps:

Step	Stock loss	DistGEMM loss	Difference
1	8.16748046875	8.16748046875	0
2	7.760293006896973	7.760293006896973	0

The raw TensorBoard grad_norm was not bitwise identical at step 2:

Step	Stock grad norm	DistGEMM grad norm	Absolute difference
1	1.7523891925811768	1.7523891925811768	0
2	1.9925340414047241	1.9925339221954346	1.1920928955078125e-7

TorchTitan's review policy requires identical loss and grad norm for a
non-computation change. The PR should either restore that property or explicitly
classify and validate this as a computation/numerics change with justified
tolerances and representative convergence evidence. A training process exiting
successfully is not numerical proof.

2. Actually execute the CUDA path in CI

The current PR checks include the 8 GPU feature and model workflows, but not the
H100 integration workflow. Their logs run features and models, not the new
h100:dist_gemm_attention entry. The CPU unit job skips
TestDistLinearPrimitives, and test_dist_gemm_attention.py explicitly does not
execute a symmetric-memory op.

The added H100 entry is useful, but it must be run before approval. It should
also prove that the intended FSDP2 + TP2 configuration was selected, rather than
only checking the process exit code.

3. Supply performance and trace evidence

This PR introduces a performance backend but reports no throughput comparison
or trace. Run at least 10 training steps and report:

exact model, batch size, sequence length, TP/DP degrees, dtype, and hardware;
warmup and throughput measurement windows;
stock versus DistGEMM throughput;
a profiler window showing the fused symmetric-memory schedules actually
selected, including whether communication/GEMM overlap occurred;
kernel/runtime counts and time, not just pass logs or operator names.

4. State and test the supported composition matrix

At minimum, cover or explicitly reject at config time:

spmd_types, default DTensor, and full_dtensor;
SP enabled and disabled;
the FSDP + TP integration added by this PR;
SAC and compile;
CUDA graph capture, because the implementation makes an explicit capture
safety guarantee.

Do not add permissive runtime fallbacks for unsupported combinations. Validate
the backend invariant before model execution with an actionable error.

tianyu-l

please migrate from dtensor to spmd_types

tianyu-l · 2026-08-12T08:24:59Z

probably should be put under models/common/

Oh I see the difference, this file is about autograd functions as building blocks.

Can we rename to distributed/linear.py as it's unambiguous?

tianyu-l · 2026-08-12T08:26:21Z

+                maybe_update_dist_gemm_config,
+            )
+
+            maybe_update_dist_gemm_config(self, config)


@fegin with your refactor let's kill this and move to config space

tianyu-l · 2026-08-12T08:28:52Z

+        # A stock rowwise linear emits a Partial over its slice of K and lets the
+        # framework reduce-scatter it. AttentionOutputLinear collapses those two
+        # steps -- the reduce-scatter happens inside the fused op -- so it returns
+        # the final Shard(1) directly and never produces a Partial. Keep only the
+        # parameter shardings: with the output already in its final layout there
+        # is nothing left to check or redistribute.


btw we'll move away from this -- we would put spmd_types redistribute collectives into the module and no longer have src->dst redistribute at module boundary, so the shardings should be the same as dist_gemm

tianyu-l · 2026-08-12T08:31:21Z

+def to_local(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor.to_local() if isinstance(tensor, DTensor) else tensor


we are deprecating dtensor backends, could you test with spmd_backend == "spmd_types" only and remove such DTensor code

tianyu-l · 2026-08-12T08:33:03Z

+    ``tokens_per_rank`` is None when the model config was never updated from a
+    runtime config (inference-only callers, unit tests). Reserving is then simply
+    skipped and the ops size the workspace lazily on first use, as they always
+    have -- correct, just without the graph-capture guarantee.


In inference, multiple cudagraphs will be captured, and different runs will use different sized cudagraph. What's the behavior?

tianyu-l · 2026-08-12T08:37:11Z

+    def forward(  # pyrefly: ignore[bad-override]
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.tp_group is None or not is_tp_sequence_sharded(x):


Are you saying that for non-SP this is not doing optimization

Yes, for non-SP, the inputs are already replicated. So, we do not need all-gather.

tianyu-l · 2026-08-12T08:38:05Z

+    return tuple(mesh_dim_names).index("tp")
+
+
+def is_tp_sequence_sharded(tensor: torch.Tensor) -> bool:


@pianpwk In spmd_types, what's the best way to test SP on / off?

tianyu-l · 2026-08-12T08:38:52Z

+        self.tp_group: dist.ProcessGroup | None = None
+        self.tokens_per_rank = config.tokens_per_rank
+
+    def parallelize(self, parallel_dims: "ParallelDims") -> None:


@fegin I think we should have a way to ban / discourage parallelize overwrite

tianyu-l · 2026-08-12T08:40:05Z

+        )
+
+
+class DistGemmGQAttention(GQAttention):


I don't think we need this class?

tianyu-l · 2026-08-12T08:42:11Z

    )


 def make_ffn_config(


How about FFN?

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/dist_linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific, so they live under `torchtitan/distributed`. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, grown to the max over layers rather than their sum, and reserved once from the model's `parallelize_fn` before any layer runs. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address may already be baked into a captured graph -- a use-after-free on replay rather than a clean error. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: ab99370 Pull-Request: #4044

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/dist_linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific, so they live under `torchtitan/distributed`. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, grown to the max over layers rather than their sum, and reserved once from the model's `parallelize_fn` before any layer runs. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address may already be baked into a captured graph -- a use-after-free on replay rather than a clean error. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: 51ccb71 Pull-Request: #4044

tianyu-l · 2026-08-13T00:40:15Z

    """
+    # Before model.parallelize(), which shards the weights this reads, and before
+    # any layer runs. A no-op unless a layer selected gemm_backend="dist_gemm".
+    reserve_dist_gemm_workspace(model, parallel_dims, training, parallelism)


We will be removing model specific parallelize functions like parallelize_llama -- could you move this to the constructors of dist gemm modules?

See https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/token_dispatcher.py#L827 for an example

Right now you still need to do the wiring of sequence length etc. into the module code, in update_from_config which @fegin will clean up later.

No longer needed. The async-TP ops themselves allocate the symmetric memory buffer, and I was over-engineering in titan for cudagraph capture. We run a few warmups before cudagraph capture anyways, so all that over-engineering was not needed.

tianyu-l · 2026-08-13T00:40:43Z

I guess we need to do this for all the models, but can be in later PRs.

tianyu-l · 2026-08-13T00:41:45Z

+    from torchtitan.distributed.parallel_dims import ParallelDims
+
+
+def tp_group_from_context() -> dist.ProcessGroup | None:


Suggested change

def tp_group_from_context() -> dist.ProcessGroup | None:

def _tp_group_from_context() -> dist.ProcessGroup | None:

also @pianpwk this looks more complicated than I expected

we can simplify after #3895 upgrades spmd_types (and just rely on mesh name "tp" + spmd_mesh_size() and maybe a helper for rank index). Though that can be a follow up, depending on when this can land

tianyu-l · 2026-08-13T00:44:35Z

Oh I see the difference, this file is about autograd functions as building blocks.

Can we rename to distributed/linear.py as it's unambiguous?

tianyu-l · 2026-08-13T00:45:17Z

+
+
+@spmd.register_autograd_function
+class AllGatherLinear(torch.autograd.Function):


do we need to custom-opify it so it's traceable and work with graph trainer? cc @sanketpurandare

No, we do not need any custom ops. Tracing works fine with autograd.function objects.

However, graph trainer does not work with spmd types as of now, so I could not add a test. I tested a per-layer compiler and that worked.

tianyu-l · 2026-08-13T00:45:47Z

+    return tp_group if tp_group.size() > 1 else None
+
+
+def local_param(param: torch.Tensor | None) -> torch.Tensor | None:


[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/models/common/dist_gemm_ops.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, grown to the max over layers rather than their sum, and reserved once from the model's `parallelize_fn` before any layer runs. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address may already be baked into a captured graph -- a use-after-free on replay rather than a clean error. The size is `2 * tokens_per_rank * max(K, N)`, which bounds every schedule these ops can pick -- but only while each `fused_all_gather_matmul` leaves `return_A` at its default. Passing False selects `_multimem_all_gather_matmul`, which reserves the full gathered buffer rather than a shard and would multiply the requirement by the TP degree. Both wgrad calls therefore leave it defaulted, which also keeps forward and backward on the same schedule. `dist_gemm_workspace_bytes` is a pure function so that arithmetic is unit-tested. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: 9d0dfcd Pull-Request: #4044

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, grown to the max over layers rather than their sum. Each module reserves its own share from `Module.parallelize` -- the generic hook every trainer calls, rather than a model-specific `parallelize_fn` a trainer might substitute (GraphTrainer does). The per-step token count it needs is stamped onto the module configs from `update_from_config`, the same way `update_ep_token_dispatcher_config` feeds the EP dispatchers. It cannot happen in the constructor -- sizing needs the TP process group, and no mesh exists at `build()` time -- nor in `forward`, where allocating is not traceable. Growing it later re-rendezvouses, which is a collective and is rejected during CUDA graph capture, and growth also frees the old buffer while its address may already be baked into a captured graph -- a use-after-free on replay rather than a clean error. The size is `2 * tokens_per_rank * max(K, N)`, which bounds every schedule these ops can pick -- but only while each `fused_all_gather_matmul` leaves `return_A` at its default. Passing False selects `_multimem_all_gather_matmul`, which reserves the full gathered buffer rather than a shard and would multiply the requirement by the TP degree. Both wgrad calls therefore leave it defaulted, which also keeps forward and backward on the same schedule. `dist_gemm_workspace_bytes` is a pure function so that arithmetic is unit-tested. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: f315bca Pull-Request: #4044

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, and the symm_mem ops size it themselves: each computes its own requirement and `get_symm_mem_workspace` grows monotonically, so the cost is a max over layers rather than a sum. Nothing here reserves it up front. Growth re-rendezvouses, which is a collective and is rejected during CUDA graph capture, but capture is always preceded by a real warmup call (`graph_trainer/cudagraph.py`), and a warmup forward and backward runs every layer's collectives -- so all growth has happened before anything is recorded. An earlier revision pre-reserved from `parallelize` to make that explicit; it bought one growth instead of a few during warmup, at the cost of duplicating the sizing arithmetic, and is not what makes capture safe. One consequence worth keeping: both wgrad all-gathers leave `return_A` at its default. Passing False would select `_multimem_all_gather_matmul`, which reserves the full gathered buffer rather than a shard -- ranks times more symmetric memory for that one call -- and its heuristic evaluates to `K <= 2048` there, a dimension it was not tuned for. Leaving it defaulted also keeps forward and backward on the same schedule. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: fca32fe Pull-Request: #4044

[ghstack-poisoned]

Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, and the symm_mem ops size it themselves: each computes its own requirement and `get_symm_mem_workspace` grows monotonically, so the cost is a max over layers rather than a sum. Nothing here reserves it up front. Growth re-rendezvouses, which is a collective and is rejected during CUDA graph capture, but capture is always preceded by a real warmup call (`graph_trainer/cudagraph.py`), and a warmup forward and backward runs every layer's collectives -- so all growth has happened before anything is recorded. An earlier revision pre-reserved from `parallelize` to make that explicit; it bought one growth instead of a few during warmup, at the cost of duplicating the sizing arithmetic, and is not what makes capture safe. One consequence worth keeping: both wgrad all-gathers leave `return_A` at its default. Passing False would select `_multimem_all_gather_matmul`, which reserves the full gathered buffer rather than a shard -- ranks times more symmetric memory for that one call -- and its heuristic evaluates to `K <= 2048` there, a dimension it was not tuned for. Leaving it defaulted also keeps forward and backward on the same schedule. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: d8e2def Pull-Request: #4044

tianyu-l

Looks neat. Please address remaining comments.

Also would love to discuss the other options in your doc tomorrow.

tianyu-l · 2026-08-14T04:44:30Z

+    ``tp_comm_overlap`` selects how the TP collectives around the QKV and output
+    projections are run. ``"none"`` leaves them to the framework, as separate
+    collectives either side of the GEMM. ``"dist_gemm"`` folds each into its
+    adjacent GEMM over symmetric memory, which requires ``fuse_qkv=True``, CUDA,


Then let's assert on fuse_qkv=True?

tianyu-l · 2026-08-14T04:46:32Z

+# next to. "none" leaves them to the framework, as separate collectives at the
+# module boundary; "dist_gemm" folds each into its adjacent GEMM over symmetric
+# memory. Named after Megatron's --tp-comm-overlap, which is the same technique.
+TpCommOverlap = Literal["none", "dist_gemm"]


--tp-comm-overlap sounds like boolean.

either make this a boolean

or let's use tp_gemm_backend with Literal options

Prefer the latter to be consistent with "attn backend", and be open about adding new solutions, like DistGEMM.

So it was earlier gemm_backend and for someone new it would sound like the options are cublas/triton/cute etc.

tp_gemm_backend is better than gemm_backend becuase it talks about tp. But, it still gives the feeling of cublas/triton/cute etc (maybe I am biased because that how I think about compiler backends).

That said I am going to switch to tp_gemm_backend to be consistent with attn_backend.

tianyu-l · 2026-08-14T04:51:21Z

+        return y_flat.view(seqlen // world_size, bsz, -1).transpose(0, 1).contiguous()
+
+
+class AllGatherFusedFeedForward(FeedForward):


Maybe let's also make this fast by default https://github.com/pytorch/torchtitan/blob/main/torchtitan/overrides/fused_swiglu.py#L429

It uses a triton silu_and_mul kernel, which is not "pytorch native" so we put it in override folder.

It also make w1 w3 in the same tensor, so I guess we could use AllGatherLinear instead of AllGatherLinearMulti?

Yeah, this can be done. Claude gave me a choice and I asked it to just focus on unfused one for now, fused one is easy to extend too.

Added in #4151 just to keep things clean

[ghstack-poisoned]

Reland of #4044 because I did a rookie mistake of merging ghstack opened PRs. Motivation Async-TP -- folding the TP all-gather and reduce-scatter into the adjacent GEMM rather than running them beside it -- is currently only reachable in torchtitan through a compiler: either torch.compile's micro-pipelining pass behind `--parallelism.enable_async_tensor_parallel`, or a GraphTrainer graph pass. That ties a parallelism strategy to a particular compilation path. This makes it a property of the model config instead, selected by a `tp_gemm_backend` argument on `model_registry` exactly as `attn_backend` already selects the attention kernel. Async-TP then works under the eager trainer with no graph pass involved. Structure `torchtitan/distributed/linear.py` holds three reusable autograd Functions. `AllGatherLinear` is a column-parallel projection (all-gather the sequence shard, then GEMM) and `LinearReduceScatter` its row-parallel dual (GEMM, then reduce-scatter). `AllGatherLinearMulti` serves several column-parallel linears that share one input, so the SwiGLU pair w1/w3 costs one all-gather rather than two; its backward still needs only the same two collectives, because the sum over outputs is expressed as a single concatenated product. None of them is attention-specific; MoE projections could use the same pair. `torchtitan/models/common/dist_gemm.py` holds the modules. `AllGatherFusedQKVLinear`, `RowParallelLinear` and `AllGatherFusedFeedForward` are drop-in replacements for the stock QKV, output and SwiGLU projections, keeping the stock parameter layouts so checkpoints still interoperate. `RowParallelLinear` serves both attention's `wo` and the FFN's `w2`. Selection follows the `attn_backend` path one for one: `model_registry` -> flavor function -> `_build_llama3_layers` -> `make_gqa_config` / `make_ffn_config`, which is where `tp_gemm_backend="dist_gemm"` picks the fused classes. Only llama3 is threaded; every other model keeps the default and can gain the argument when someone needs it. `llama3_debugmodel_dist_gemm` is the config-registry entry, and the h100 integration test runs it at TP=2. No attention or FFN subclass is needed beyond the projections themselves. That required one fix in shared code: `GQAttention.forward` captured `B, L` from its input and reused them to fold the attention output, which is only valid when the QKV preserves sequence length. `AllGatherFusedQKVLinear` gathers the SP shard, so it does not, and under spmd_types -- where a tensor's shape is its local shape -- that silently folded sequence into features. It now derives the fold from the tensor being folded. spmd_types only The fused modules take and return plain local tensors, which is the `spmd_types` contract; the DTensor backends are being deprecated and are rejected outright. Two preconditions are checked at parallelize time, because neither is detectable from inside a module once activations carry no placements: the backend must be `spmd_types`, and sequence parallelism must be enabled. The fused GEMMs *are* the SP collectives, so with SP disabled there is nothing to gather and `wo` would reduce-scatter where it must all-reduce. Because selection happens at config-construction time, `set_gqa_attention_sharding` and `set_dense_ffn_sharding` declare the fused blocks' contracts directly: no boundary all-gather, and a row-parallel output that emits its final Shard(1) rather than the Partial a stock rowwise linear produces. Both branches are transitional and collapse once redistribute collectives move inside the modules generally. Symmetric memory One workspace per process group serves every layer, and the symm_mem ops size it themselves: each computes its own requirement and `get_symm_mem_workspace` grows monotonically, so the cost is a max over layers rather than a sum. Nothing here reserves it up front. Growth re-rendezvouses, which is a collective and is rejected during CUDA graph capture, but capture is always preceded by a real warmup call (`graph_trainer/cudagraph.py`), and a warmup forward and backward runs every layer's collectives -- so all growth has happened before anything is recorded. An earlier revision pre-reserved from `parallelize` to make that explicit; it bought one growth instead of a few during warmup, at the cost of duplicating the sizing arithmetic, and is not what makes capture safe. One consequence worth keeping: both wgrad all-gathers leave `return_A` at its default. Passing False would select `_multimem_all_gather_matmul`, which reserves the full gathered buffer rather than a shard -- ranks times more symmetric memory for that one call -- and its heuristic evaluates to `K <= 2048` there, a dimension it was not tuned for. Leaving it defaulted also keeps forward and backward on the same schedule. Every op carves its buffers from offset 0 of that shared workspace, so the regions alias; safety comes from the ops being serialized rather than from disjointness. Each op also brackets itself with barriers, so ranks cannot run ahead of one another. Sequential module forwards and autograd backward are single-stream, which is how models actually run, so this holds today -- but deliberate cross-stream overlap would need distinct workspace offsets, not a barrier. Known gap `--debug.spmd_typechecking` does not pass. The autograd Functions are registered with `typecheck_forward` declaring their type transitions, which the checker requires, but it then rejects the QKV reshape that splits the TP-sharded feature dim into (n_kv, r_dim, head_dim) -- aligned in practice, since the shard is over n_kv. The stock `FusedQKVLinear` hits the same limitation and works around it with `spmd.local()` plus an explicit `PartitionSpec`, carrying a TODO to remove that once the checker handles aligned splits. Rather than add a second instance of a pattern already marked for deletion, this leaves typechecking unsupported until that lands upstream. Authored with assistance from Claude Code. ghstack-source-id: e001593 Pull-Request: #4044

Update

49b87d3

[ghstack-poisoned]

pytorch-bot Bot added the ciflow/8gpu label Jul 31, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 31, 2026

anijain2305 changed the title ~~Add eager dist GEMM attention override~~ [wip][not for review] Add eager dist GEMM attention override Jul 31, 2026

Update

c1b2543

[ghstack-poisoned]

anijain2305 added a commit that referenced this pull request Aug 6, 2026

Add eager dist GEMM attention override

7bb5afa

ghstack-source-id: debdd0f Pull-Request: #4044

Update

f93a01c

[ghstack-poisoned]

anijain2305 added a commit that referenced this pull request Aug 6, 2026

Add eager dist GEMM attention override

22d1cfd

ghstack-source-id: aee34d3 Pull-Request: #4044

Update

013e578

[ghstack-poisoned]

anijain2305 changed the title ~~[wip][not for review] Add eager dist GEMM attention override~~ Add eager dist GEMM attention override Aug 10, 2026

Update

511092b

[ghstack-poisoned]

Update

ac7708b

[ghstack-poisoned]

Update

778f856

[ghstack-poisoned]

Update

8665dc5

[ghstack-poisoned]

tianyu-l reviewed Aug 11, 2026

View reviewed changes

Update

6aa603a

[ghstack-poisoned]

anijain2305 requested a review from sanketpurandare August 11, 2026 21:15

anijain2305 marked this pull request as ready for review August 11, 2026 21:15

anijain2305 requested review from fegin, wconstab and wwwjn as code owners August 11, 2026 21:15

Update

a6a0549

[ghstack-poisoned]

anijain2305 changed the title ~~Add eager dist GEMM attention override~~ Add dist GEMM attention backend Aug 11, 2026

sanketpurandare reviewed Aug 12, 2026

View reviewed changes

tianyu-l requested changes Aug 12, 2026

View reviewed changes

Update

96561e0

[ghstack-poisoned]

anijain2305 changed the title ~~Add dist GEMM attention backend~~ Add dist GEMM attention and FFN backend Aug 12, 2026

Update

aa8e4c8

[ghstack-poisoned]

tianyu-l reviewed Aug 13, 2026

View reviewed changes

Update

2082ac2

[ghstack-poisoned]

Update

846c09c

[ghstack-poisoned]

anijain2305 requested review from IvanKobzarev, SherlockNoMad, aditvenk and xmfan as code owners August 13, 2026 23:52

Update

abe4242

[ghstack-poisoned]

Update

89bd07e

[ghstack-poisoned]

anijain2305 added the ciflow/h100.8 Trigger H100.8 CI label Aug 14, 2026

tianyu-l approved these changes Aug 14, 2026

View reviewed changes

tianyu-l mentioned this pull request Aug 14, 2026

[New Model]Add DeepSeek V4 model support #3634

Draft

Update

cb0ab88

[ghstack-poisoned]

anijain2305 mentioned this pull request Aug 14, 2026

Support the fused SwiGLU override with dist-GEMM TP overlap #4151

Open

anijain2305 merged commit 7370523 into gh/anijain2305/3/base Aug 14, 2026
13 of 14 checks passed

anijain2305 mentioned this pull request Aug 15, 2026

[Reland] Add dist GEMM attention and FFN backend #4162

Merged

		def to_local(tensor: torch.Tensor) -> torch.Tensor:
		return tensor.to_local() if isinstance(tensor, DTensor) else tensor

		return tuple(mesh_dim_names).index("tp")


		def is_tp_sequence_sharded(tensor: torch.Tensor) -> bool:

		from torchtitan.distributed.parallel_dims import ParallelDims


		def tp_group_from_context() -> dist.ProcessGroup \| None:

	def tp_group_from_context() -> dist.ProcessGroup \| None:
	def _tp_group_from_context() -> dist.ProcessGroup \| None:



		@spmd.register_autograd_function
		class AllGatherLinear(torch.autograd.Function):

		return tp_group if tp_group.size() > 1 else None


		def local_param(param: torch.Tensor \| None) -> torch.Tensor \| None:

		return y_flat.view(seqlen // world_size, bsz, -1).transpose(0, 1).contiguous()


		class AllGatherFusedFeedForward(FeedForward):

Uh oh!

Conversation

anijain2305 commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sanketpurandare commented Aug 12, 2026

1. Resolve or characterize the numerical drift

2. Actually execute the CUDA path in CI

3. Supply performance and trace evidence

4. State and test the supported composition matrix

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pianpwk Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anijain2305 Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

anijain2305 commented Jul 31, 2026 •

edited

Loading

pianpwk Aug 14, 2026 •

edited

Loading

anijain2305 Aug 14, 2026 •

edited

Loading