[BC][MLA][CP] Optimize MLA all-gather and Ulysses communication for Kimi K3 and dpskV3 by acisseJZhong · Pull Request #4760 · pytorch/torchtitan

acisseJZhong · 2026-09-17T00:52:10Z

Stack order: #5445 -> #4760.

Naming and file moves are in #5445; this PR focuses on MLA CP optimization.

The original optimization idea was inspired by this MLA Ulysses CP proposal.

Shape suffixes: T tokens, H heads, L KV latent, N non-positional key, R shared key, V value; K=N+R, P=N+V, F=K+P. Suffixes describe logical dimensions; collectives change their sharding.

K/V all-gather CP

Before: all-gather materialized K and V separately. K repeats k_shared_TR
across attention heads.
After: all-gather normalized compressed KV latents and head-shared
k_shared_TR, then apply wkv_b and materialize K/V locally. This trades
repeated projection work for lower communication volume.

# Before: materialize K before communication.
kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL))
k_nope_THN, v_THV = unpack(kv_THP)
k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR))
k_THK = all_gather(k_THK)
v_THV = all_gather(v_THV)

# After: communicate compressed latents before projecting K/V.
kv_c_normed_TL = all_gather(kv_c_normed_TL)
k_shared_TR = all_gather(k_shared_TR)
kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL))
k_nope_THN, v_THV = unpack(kv_THP)
k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR))

Ulysses CP

Before: use three all-to-alls for materialized Q, K, and V. K repeats
k_shared_TR across attention heads.
Current: all-gather headless k_shared_TR, then use one all-to-all
for q_THK + packed kv_THP.

# Before: materialize K, then move Q, K, and V from token to head sharding.
kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL))
k_nope_THN, v_THV = unpack(kv_THP)
k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR))
q_THK = all_to_all_token_to_head(q_THK)
k_THK = all_to_all_token_to_head(k_THK)
v_THV = all_to_all_token_to_head(v_THV)
out_THV = attention(q_THK, k_THK, v_THV)
out_THV = all_to_all_head_to_token(out_THV)

# Current: combine q_THK and kv_THP in one all-to-all.
kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL))
k_shared_TR = all_gather(k_shared_TR)
packed_q_kv_THF = all_to_all_token_to_head(
    concat(q_THK, kv_THP)
)
q_THK, kv_THP = split(packed_q_kv_THF)
k_nope_THN, v_THV = unpack(kv_THP)
k_THK = concat(
    k_nope_THN,
    expand_per_head(k_shared_TR),
)
out_THV = attention(q_THK, k_THK, v_THV)
out_THV = all_to_all_head_to_token(out_THV)

Shared MLA backend and checkpoints
- DeepSeek V3, Kimi K2.7, and Kimi K3 share the compact MLA backend.
  MLAAttention owns wkv_b and delegates to a Q/K/V attention kernel;
  common all-gather communication lives in distributed/context_parallel.py.
- HF adapters map directly to attention.inner_attention.wkv_b.weight,
  replacing state-dict renaming hooks. HF keys remain compatible. Native
  checkpoints use the nested key; old attention.wkv_b.weight keys are no
  longer automatically remapped.
Sharding and backend compatibility
- The KV projection explicitly asserts its CP layout: replicated after latent
  all-gather, token-sharded before Ulysses all-to-all. The CP transform copies
  the nested projection config when selecting the backend; the local head
  reshape preserves that layout.
- fused_mla rejects CP backends at model construction because its fused
  forward bypasses the compact MLA communication wrapper. Fused MLA with CP
  is not supported; remove the override to use the configured CP backend.
- GraphTrainer precompile delegates metadata construction to the Decoder;
  batch-invariant settings traverse Flex configs directly, including those
  nested inside MLA.
- The TorchTitan vLLM wrapper currently rejects MLA. Kimi K3 MLA inference
  parity and RL integration tests are disabled until support is restored.

Tests

After rebase/split at 4d5064c40c: 197 CPU tests and the four-GPU MLA sharding test passed. Non-Pyrefly lint passed; Pyrefly has the same 53 diagnostics as main. GraphTrainer validation remains blocked by the installed PyTorch missing _is_cacheable_effect.

Numerical comparison

3e6c5c2b1a: DeepSeek V3 and text-only Kimi K3 debug models, 2 H100s, CP=2, BF16/Flex, seed 42, deterministic, 256 tokens, 100 steps. Identical initial weights and a fixed synthetic batch; AdamW LR 1e-3, gradient clipping 1.0.

Relative differences between the materialized-K/V reference and compact MLA on the same revision. Loss is global mean cross-entropy; grad norm is FP32 L2 before clipping.

Model	CP backend	Loss Δ, step 10	Grad norm Δ, step 10	Loss Δ, step 100	Grad norm Δ, step 100
DeepSeek V3	K/V all-gather	0.067%	0.371%	0.541%	0.116%
DeepSeek V3	Ulysses	0.103%	0.218%	0.056%	0.128%
Kimi K3	K/V all-gather	2.321%	11.125%	0.462%	0.628%
Kimi K3	Ulysses	0.262%	2.466%	0.363%	0.322%

All runs completed 100 steps with finite loss and gradients. First-step logits and loss are bitwise identical in all four comparisons, but gradients are not: first-step relative L2 gradient differences range from 0.1475% to 0.2716%. Changed BF16 gradient accumulation leads to different subsequent optimizer updates. This fixed-batch comparison does not establish convergence on a training dataset.

Performance

Retested at 3e6c5c2b1a on MAST: DeepSeek-V3 16B, 32 H100 GPUs, CP=8/FSDP=4/EP=8, BF16 mixed precision, Flex attention, C4 test data, CUDA graphs disabled, 100 steps. PyTorch 2.15.0a0+git8646abd.

The baseline materializes K/V before communication on the same revision and runtime. TPS/GPU is the rank-0 mean +/- step SD over steps 21-99, from logged throughput. One run per case.

K/V all-gather CP

32,768 tokens per DP rank; HeadTail load balancing.

Path	Input collectives	TPS/GPU (mean +/- SD)	vs. baseline
Before: materialized K/V	2 all-gathers	3817.5 +/- 1021.8	-
Current: compressed KV latent + shared key	2 all-gathers	3970.8 +/- 1009.7	+4.02%

Ulysses CP

8,192 tokens per DP rank; no CP load balancing.

Path	Input collectives	TPS/GPU (mean +/- SD)	vs. baseline
Before: separate Q, K, V	3 all-to-alls	1000.6 +/- 289.5	-
Current: combined Q + packed KV, shared key	1 all-to-all + 1 all-gather	1098.7 +/- 303.9	+9.80%

Archived notes.

fegin

Put MLAAttention in common if you want to preserve _materialize_kv approach.

fegin · 2026-09-22T06:38:46Z

+# P = packed per-head KV channels (N + V)
+
+
+class MLAInnerAttention(InnerAttention, ABC):


Why putting this for Kimi_k3 only? DSv3 can actually share this. I don't see a difference.

acisseJZhong · 2026-10-07T06:14:00Z

+        **kwargs,
+    ) -> torch.Tensor:
+        k_THK, v_THV = materialize_mla_kv(q_THK, kv_THP, k_shared_TR)
+        return VLLMInnerAttention.forward(


this defeats the purpose of mla as it still caches full kv. Left a todo for now.

fegin · 2026-10-07T20:14:57Z

        KVAllGatherCPSlidingWindowFlexInnerAttention,
    )
+    from torchtitan.models.common.attention.cp_mla import (
+        KVAllGatherCPMLAFlexInnerAttention,


We should remove "CP" prefix:

The naming is already too long.

KVAllGather and Ulysses already show CP meaning.

CP as upper cases make the reading much harder.

We don't need to do that in this PR, but I feel this is needed. My PR has KVAllGatherCPVarlenInnerAttention and KVAllGatherCPSlidingWindowVarlenInnerAttention. It is hard to read.

fegin · 2026-10-07T20:15:58Z

+from torchtitan.models.common.attention import (
+    materialize_mla_kv,
+    MLAInnerAttention,
+    ScaledDotProductInnerAttention,


Any reason why graph trainer needs ScaledDotProductInnerAttention? cc., @SherlockNoMad

fegin · 2026-10-07T20:23:31Z

+    k_shared_TR = spmd.redistribute(
+        k_shared_TR,
+        cp_group,
+        src=spmd.S(_TOKEN_DIM),
+        dst=spmd.R,
+        backward_options={"op_dtype": k_shared_TR.dtype},
+    )


One minor improvement:
Move this redistribute before packed_q_kv_THF cat. This way we may be able to gain some overlapping, you probably see no improvement at all since cat should be fast, but this change cost nothing.

acisseJZhong · 2026-10-13T23:32:00Z

    class Config(Module.Config):
        n_heads: int
-        inner_attention: InnerAttention.Config
+        inner_attention: "InnerAttention.Config | MLAAttention.Config"


this may look weird, codex proposes an alternative way to use

class Attention(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): n_heads: int backend: "InnerAttention.Config | MLAAttention.Config" class MLAAttention(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): wkv_b: Linear.Config packed_kv_head_dim: int inner_attention: InnerAttention.Config

I suggest not to mess up this class with MLA

tianyu-l · 2026-10-14T06:27:01Z

-                inner.kernel_options["BLOCK_M"] = self._BLOCK_M
-                inner.kernel_options["BLOCK_N"] = self._BLOCK_N
+            inner_attention_config = (
+                layer_cfg.attention.inner_attention.inner_attention_config


this looks not right

tianyu-l · 2026-10-14T06:28:18Z

+    inner_attention_map: Mapping[
+        type[InnerAttention | MLAAttention], type[InnerAttention | MLAAttention]
+    ]


wait, it'd be Mapping[type Inner, type Inner] | Mapping[type MLA, type MLA]?

tianyu-l · 2026-10-14T06:32:07Z

    FlexAttentionMetadata | VarlenAttentionMetadata | LinearAttentionMetadata
 )
-AttentionMetadataMap = Mapping[type[InnerAttention], AttentionMetadata]
+AttentionMetadataMap = Mapping[type[InnerAttention | MLAAttention], AttentionMetadata]


is MLAAttention outer or inner or somewhere in the middle? hard to tell from naming

did you end up choosing the 4 layer structure?

this PR is mixing two optimizations?

maybe let's chat in person.

tianyu-l · 2026-10-14T06:32:26Z

    @dataclass(kw_only=True, slots=True)
    class Config(Module.Config):
+        @property
+        def inner_attention_config(self) -> "InnerAttention.Config":


removed this field and switched to get_attention_kernel_config. This is used to extract the inner inner attention for MLA.

tianyu-l · 2026-10-14T06:33:26Z

    class Config(Module.Config):
        n_heads: int
-        inner_attention: InnerAttention.Config
+        inner_attention: "InnerAttention.Config | MLAAttention.Config"


I suggest not to mess up this class with MLA

tianyu-l · 2026-10-14T06:34:13Z


 class ContextParallelInnerKDA(
-    CPInnerAttention[
+    CPAttention[


hmm what happened

tianyu-l · 2026-10-14T06:34:42Z

i think we need a chart for MLA <> CP <> inference

we could add once after inference story is completed.

tianyu-l · 2026-10-16T00:33:00Z

        return out_TD.view(num_tokens, -1, self.value_head_dim)


+class VLLMMLAAttention(MLAAttention):


maybe put in vllm/mla.py as it's going to have its own complexity

vllm mla is banned for now, so this class is deleted.

tianyu-l · 2026-10-16T00:33:34Z


-class CPInnerAttention(
-    InnerAttention,
+class CPAttention(


I didn't get this change.

If every module that inherits CPAttention is also inheriting InnerAttention, should we make CPAttention inheriting InnerAttention?

tianyu-l · 2026-10-16T00:33:52Z

+        return self.inner_attention(q_THK, k_THK, v_THV, **kwargs)
+
+
+def get_attention_kernel_config(


not clear to me why we need this yet

tianyu-l · 2026-10-16T01:02:51Z

+                ),
+                "sharding_config": inner_attention_config.sharding_config,
+            }
+            if isinstance(attention_backend_config, MLAAttention.Config):


let's ban this for now, so that you don't need get_attention_kernel_config

acisseJZhong · 2026-10-16T01:03:59Z

@claude review @tianyu-l 's latest comments

claude · 2026-10-16T01:04:21Z

Claude finished @acisseJZhong's task in 2m 14s —— View job

Review of @tianyu-l's latest comments (2026-10-16)

Gather context (mla.py, cp_attention.py, cp_mla.py, vllm/attention.py, model_wrapper.py, every get_attention_kernel_config callsite)
Go through each of the 4 latest comments
Post recommendations

Summary: I agree with 3 of the 4 comments as written. I'd push back on the 4th, banning MLA in vLLM, because Kimi K3 RL recipes already exist on main. There's another way to drop get_attention_kernel_config that keeps vLLM working.

1. `cp_attention.py:56`: "If every module that inherits CPAttention also inherits InnerAttention, should CPAttention inherit InnerAttention?"

Agree. Every concrete subclass already inherits InnerAttention:

CP class	non-CP base	Is it an `InnerAttention`?
`KVAllGatherFlexInnerAttention`, `Ulysses*InnerAttention`	`FlexInnerAttention` / `VarlenInnerAttention`	yes
`KVAllGatherMLAFlexAttention`, `UlyssesMLAAttention`	`MLAAttention`	yes (`MLAAttention(InnerAttention)`, `mla.py:52`)
`ContextParallelInnerKDA`	`InnerKDA`	yes (`kda.py:141`)
`ContextParallelInnerGDN`	`InnerGDN`	yes (`gdn.py:268`)

So the switch from CPInnerAttention(InnerAttention, ...) to CPAttention(Module, ...) gains nothing. It also causes the unrelated churn in cp_kda.py / cp_gdn.py / validation.py / decoder.py that the earlier "hmm what happened" comment asked about. My suggestion is to restore main's shape and name:

class CPInnerAttention(
    InnerAttention, ABC, Generic[_GlobalAttentionMetadataT, _LocalAttentionMetadataT]
):
    @dataclass(kw_only=True, slots=True)
    class Config(InnerAttention.Config):
        pass

The MRO should still resolve, e.g. UlyssesMLAAttention -> Ulysses -> CPInnerAttention -> MLAAttention -> InnerAttention -> Module. That's the same diamond main already has for KVAllGatherFlexInnerAttention(KVAllGatherFlex, FlexInnerAttention). Both intermediate Configs add no slots, so slots=True won't hit a layout conflict. (I couldn't run Python in this job, so please confirm locally.) This also settles the earlier Mapping[type Inner, type Inner] | Mapping[type MLA, type MLA] question: ContextParallelTransform.inner_attention_map can stay Mapping[type[InnerAttention], type[InnerAttention]], as it already is at context_parallel.py:33. fegin's "drop the CP prefix" rename can come in a follow-up.

Fix this →

2. `vllm/attention.py:453`: "maybe put in vllm/mla.py"

VLLMMLAAttention is currently an empty subclass (Config(MLAAttention.Config): pass, no behavior). I'd delete it rather than move it. _replace_vllm_layer_configs only needs to swap the kernel inside the existing MLA config:

if isinstance(backend_cfg, MLAAttention.Config):
    vllm_attention_cfg = dataclasses.replace(
        backend_cfg, inner_attention=VLLMInnerAttention.Config(**config_kwargs)
    )

Once the native vLLM MLA backend arrives (the TODO about caching [kv_c_normed, k_pe]), it will have real logic, and that's when vllm/mla.py makes sense. Creating the file now would only hold an empty class. You'd also need to update tests/rl/unit_tests/cpu/test_vllm_wrapper.py (lines 24, 210, 219, 240, 257).

3 & 4. `mla.py:113`: "not clear why we need this" and `model_wrapper.py:102`: "let's ban this for now, so that you don't need `get_attention_kernel_config`"

I'd push back on banning MLA in vLLM.

These Kimi K3 RL recipes already exist on main and run MLA through vLLM: rl_grpo_kimi_k3_debug_varlen, rl_grpo_kimi_k3_debug_varlen_batch_invariant and rl_grpo_kimi_k3_debug_batch_invariant in torchtitan_recipes/tests/rl/alphabet_sort.py.
tests/rl/unit_tests/gpu/test_bitwise_parity.py:929 depends on the batch-invariant one.
Banning would turn this CP PR into an RL regression. That violates principle 5 ("Protect battle-tested code paths").

Banning also wouldn't remove the helper. It has 5 callsites, and only 3 of them are vLLM-related:

Callsite	Why it needs the inner kernel	Removed by the ban?
`rl/generator/vllm/model_wrapper.py:82`	`sharding_config`	yes
`rl/generator/generator.py:969`	Flex/Varlen check	yes
`rl/generate.py:133`	Flex/Varlen check	yes
`config/transform/batch_invariance.py:43`	sets Flex `kernel_options` for training	no
`experiments/graph_trainer/precompile_main.py:234`	decides whether to build attention metadata	no

Suggested alternative: drop get_attention_kernel_config in favor of the existing Configurable.traverse. MLAAttention.Config holds its kernel as an ordinary nested field, so the generic config walk (config/configurable.py:77, already used for AuxLoss.Config in precompile_main.py:254) finds it without any MLA-specific helper:

# batch_invariance.py
for _, flex_cfg, _, _ in layer_cfg.attention.traverse(FlexInnerAttention.Config):
    flex_cfg.kernel_options.update(BACKEND="TRITON", BLOCK_M=..., BLOCK_N=...)

# generator.py / generate.py / precompile_main.py
kernel_types = (VarlenInnerAttention.Config, FlexInnerAttention.Config)
has_supported_kernel = next(attention.traverse(kernel_types), None) is not None

traverse takes type[_ConfigT], so a tuple might need a small widening or two traverse calls.

model_wrapper.py then only needs sharding_config. With fix 2, the vLLM kernel's sharding can be read from the replaced kernel directly (backend_cfg.inner_attention.sharding_config inside the MLA branch, attention_cfg.inner_attention.sharding_config otherwise). That's a single local isinstance, not a public helper.

After that, get_attention_kernel_config can come out of mla.py and __init__.py. This avoids both the ban and the MLA-aware helper leaking into config transforms and experiments.

Fix this →

Related, optional

config/transform/context_parallel.py:72-88 special-cases the two cp_mla classes so it can rewrite wkv_b's sharding. That's MLA knowledge inside a generic transform, much like the batch-invariance concern. It could move into KVAllGatherMLAFlexAttention.Config.__post_init__ / UlyssesMLAAttention.Config.__post_init__, e.g. self.wkv_b = replace(self.wkv_b, sharding_config=mla_kv_projection_sharding_config(cp=...)). The transform would then stay a plain type swap.

I didn't change any code (review only), and Python and the tests weren't run in this job.

fegin

I suggest that split the renaming part to another PR. This PR should focus on MLA CP optimization.
You need to rebase and address muse spark as well.

tianyu-l · 2026-10-16T21:56:17Z

+            # MLA attention is not supported by the TorchTitan vLLM wrapper yet.
+            disabled=True,


cc @liangel-02

tianyu-l · 2026-10-16T22:00:02Z

+        # The fused forward bypasses the MLA backend that owns CP communication.
+        if isinstance(config.inner_attention, CPInnerAttention.Config):
+            raise ValueError(
+                "FusedMLAAttention does not support context parallelism. "
+                "Remove the fused_mla override to use the configured CP attention backend."
+            )


OK to unblock, but what's our long-term plan for this? Giving up fused MLA during CP sounds leaving room for improvement.

I can follow up in another PR

pytorch-bot · 2026-10-17T21:55:50Z

Warning: Unknown label ciflow/8gpu.
Currently recognized labels are

ciflow/b200
ciflow/docker
ciflow/fake-pg
ciflow/h100.8
ciflow/real-pg
ciflow/rl
ciflow/rocm

Please add the new label to .github/pytorch-probot.yml

fegin · 2026-10-19T17:20:05Z

 class AllGatherKVSlidingWindowFlexInnerAttention(
-    AllGatherKVFlex, SlidingWindowFlexInnerAttention
+    AllGatherKVFlex,
+    SlidingWindowFlexInnerAttention,


nit, why an extra line now?

fegin · 2026-10-19T18:01:55Z

        output = remat.region(
            self.inner_attention,
            self.remat_region_name("inner_attention"),
            recompute=self.remat_should_recompute("inner_attention"),


ActivationRematConfig( save_regions=["attention.inner_attention"], )

This doesn't work anymore. The root cause, wkv_b, is a linear, another remat region. If we are going to save "attention.inner_attention", we must save "attention.inner_attention.wkv_b", otherwise, we will get a runtime error.

One solution is to remove the remat.region here and move it into MLAAttention and separately wrap the kernel and cp_region. kernel recompute follows attention.inner_attention and cp_region recompute follows kernel (hence follow attention.inner_attention).

fegin

Also, this PR also break checkpoint BC, we should note in title [BC]. This is bigger than code BC if some people are already training dsv3.

fegin · 2026-10-19T19:48:49Z

+                "CP attention requires an active multi-rank CP mesh axis."
+            )
+
+        kv_THP = self.project_mla_kv(kv_c_normed_TL)


Move this after the redistribute below, this give us some chance to overlap the comm and comp.

fegin

one final change, thanks

fegin · 2026-10-19T20:35:37Z

+            k_shared_TR,
+            cp_group,
+            src=spmd.S(_TOKEN_DIM),
+            dst=spmd.R,
+            backward_options={"op_dtype": k_shared_TR.dtype},


Add async_op=True. https://github.com/meta-pytorch/spmd_types/blob/main/spmd_types/_collectives.py#L1239

acisseJZhong added this pull request to stack #4640 September 17, 2026 00:52

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

acisseJZhong force-pushed the optim_ulyssescp branch 7 times, most recently from 83be25f to 2a83004 Compare September 17, 2026 07:56

acisseJZhong changed the title ~~[do not review][CP]Optimize Ulysses Comm for MLA~~ [CP][Kimi K3] Add communication efficient MLA CP backends Sep 17, 2026

acisseJZhong changed the title ~~[CP][Kimi K3] Add communication efficient MLA CP backends~~ [CP]Optimize Kimi K3 MLA all-gather and Ulysses communication Sep 17, 2026

acisseJZhong marked this pull request as ready for review September 17, 2026 08:14

acisseJZhong requested review from fegin, tianyu-l, wconstab and wwwjn as code owners September 17, 2026 08:14

pytorch-bot Bot added the ciflow/fake-pg Run 1-GPU Fake PG integration tests label Sep 17, 2026

acisseJZhong changed the title ~~[CP]Optimize Kimi K3 MLA all-gather and Ulysses communication~~ [CP] Optimize Kimi K3 MLA all-gather and Ulysses communication Sep 17, 2026

acisseJZhong force-pushed the optim_ulyssescp branch 3 times, most recently from 7141d4b to 69f2ea2 Compare September 17, 2026 22:11

pytorch-bot Bot added the ciflow/rl label Sep 17, 2026

acisseJZhong force-pushed the optim_ulyssescp branch 2 times, most recently from a7497ea to b2d187c Compare September 21, 2026 06:22

acisseJZhong force-pushed the optim_ulyssescp branch from b2d187c to 6e40f0e Compare September 21, 2026 17:44

fegin requested changes Sep 22, 2026

View reviewed changes

acisseJZhong force-pushed the optim_ulyssescp branch from 6e40f0e to e96aa78 Compare September 23, 2026 02:11

acisseJZhong force-pushed the optim_ulyssescp branch 2 times, most recently from 452a5d9 to 0b1b18f Compare October 7, 2026 01:46

acisseJZhong marked this pull request as ready for review October 7, 2026 02:06

acisseJZhong requested review from IvanKobzarev, SherlockNoMad, aditvenk and sanketpurandare as code owners October 7, 2026 02:06

pytorch-bot Bot added the ciflow/fake-pg Run 1-GPU Fake PG integration tests label Oct 7, 2026

tianyu-l requested a review from liangel-02 October 7, 2026 02:25

acisseJZhong requested a review from felipemello1 as a code owner October 7, 2026 06:10

acisseJZhong commented Oct 7, 2026

View reviewed changes

fegin reviewed Oct 7, 2026

View reviewed changes

This was referenced Oct 11, 2026

Remove shared SDPA and use Varlen attention in Flux #5295

Merged

[Attn] Remove SDPA and use Varlen attention in Flux #5296

Merged

acisseJZhong commented Oct 14, 2026

View reviewed changes

tianyu-l requested changes Oct 14, 2026

View reviewed changes

tianyu-l reviewed Oct 16, 2026

View reviewed changes

fegin reviewed Oct 16, 2026

View reviewed changes

acisseJZhong mentioned this pull request Oct 16, 2026

[Attn] Rename CP backends and relocate model attention modules #5445

Merged

tianyu-l approved these changes Oct 16, 2026

View reviewed changes

fegin requested changes Oct 19, 2026

View reviewed changes

fegin reviewed Oct 19, 2026

View reviewed changes

[CP] Optimize MLA all-gather and Ulysses communication

5b2b116

fegin reviewed Oct 19, 2026

View reviewed changes

fegin approved these changes Oct 19, 2026

View reviewed changes

[MLA] Support selective rematerialization

5018b40

acisseJZhong mentioned this pull request Oct 20, 2026

[CI Fix][ROCm] Refresh DeepSeek V3 CP MI350X golden #5557

Merged

		# P = packed per-head KV channels (N + V)


		class MLAInnerAttention(InnerAttention, ABC):

		return out_TD.view(num_tokens, -1, self.value_head_dim)


		class VLLMMLAAttention(MLAAttention):

		return self.inner_attention(q_THK, k_THK, v_THV, **kwargs)


		def get_attention_kernel_config(

		# MLA attention is not supported by the TorchTitan vLLM wrapper yet.
		disabled=True,

Conversation

acisseJZhong commented Sep 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tests

Numerical comparison

Performance

K/V all-gather CP

Ulysses CP

Uh oh!

fegin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

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!

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!

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!

acisseJZhong commented Oct 16, 2026

Uh oh!

claude Bot commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review of @tianyu-l's latest comments (2026-10-16)

1. cp_attention.py:56: "If every module that inherits CPAttention also inherits InnerAttention, should CPAttention inherit InnerAttention?"

2. vllm/attention.py:453: "maybe put in vllm/mla.py"

3 & 4. mla.py:113: "not clear why we need this" and model_wrapper.py:102: "let's ban this for now, so that you don't need get_attention_kernel_config"

Related, optional

Uh oh!

fegin 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!

pytorch-bot Bot commented Oct 17, 2026

acisseJZhong commented Sep 17, 2026 •

edited

Loading

claude Bot commented Oct 16, 2026 •

edited

Loading

1. `cp_attention.py:56`: "If every module that inherits CPAttention also inherits InnerAttention, should CPAttention inherit InnerAttention?"

2. `vllm/attention.py:453`: "maybe put in vllm/mla.py"

3 & 4. `mla.py:113`: "not clear why we need this" and `model_wrapper.py:102`: "let's ban this for now, so that you don't need `get_attention_kernel_config`"