[New Model]Add DeepSeek V4 model support by Matrix-Z97 · Pull Request #3634 · pytorch/torchtitan

Matrix-Z97 · 2026-06-11T12:48:12Z

Description

This PR adds deepseek_v4 model support and registers it in TorchTitan’s supported model list.

Main Changes

Add the new torchtitan/models/deepseek_v4/ model directory
Implement the DeepSeek V4 model architecture, sparse attention, compressor, MoE, mHC, sharding, parallelization, and state dict adapter
Add a deepseek_v4 debug model configuration
Register deepseek_v4 in torchtitan/models/init.py

Parallelism Support

FSDP is currently supported
EP (Expert Parallel) is currently supported
TP (Tensor Parallel) is currently supported
CP (Context Parallel) is not supported yet and explicitly raises NotImplementedError

Additional Notes

The attention implementation is currently based on small operators rather than a fused/custom kernel.
Only the debug model configuration is included for now.

If needed, I can further improve and complete this implementation.

pytorch-bot · 2026-06-11T12:48:19Z

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

Matrix-Z97 · 2026-06-12T08:52:34Z

@tianyu-l I noticed that you added deepseek_v4 to the roadmap. However, I'm sorry that I just force-pushed a new version of the code, which caused the roadmap content to disappear. Could you please add it back?

tianyu-l · 2026-06-12T08:56:50Z

@Matrix-Z97 you mean the "26H1 TorchTitan Development"?

Oh that's not a roadmap but just for tracking what this repo is doing. I re-marked the status from "Done" to "In Progress".

Matrix-Z97 · 2026-06-12T09:03:32Z

@tianyu-l Understood, thanks for the correction! Out of curiosity, are you interested in the deepseek_v4 model?

tianyu-l · 2026-06-12T09:10:09Z

Yes, of course. We have plans, but were lacking bandwidth.

This PR comes in very timely, and the quality looks very high. We just need some time to review. Many thanks!

Btw @shuhuayu is working on adding latest Kimi model.

Matrix-Z97 · 2026-06-15T08:46:34Z

Thanks a lot for the feedback! Happy to help with the bandwidth. Let me know if there's anything I can do to help move this PR forward or if you need any adjustments. And awesome news about the latest Kimi model @shuhuayu.

tianyu-l · 2026-06-17T00:24:35Z

cc @drisspg to review the attention part

@tianyu-l How do you feel about the current FlexAttention-based DSV4 attention implementation?

It relies on the passed-in topk_idxs to construct the attention mask. However, because the resulting mask is quite sparse, the attention computation performance is not ideal. Do you have any suggestions here? For example, would it be better to use a custom sparse attention kernel based on TileLang, Triton, Helion, or another approach?

I think the major concern is that flex is optimized for block masks, while topk_idxs is a token-wise selection, so scattered picks leave few empty blocks for flex to skip. Since the selection is untrained for now, I'd suggest we first land the aux loss to train it, then revisit this optimization.

floatingtrees · 2026-06-29T20:11:04Z

@Matrix-Z97 Are you planning to add MTP support for this model? If not, I’d be happy to add it.

Matrix-Z97 · 2026-07-02T02:15:48Z

@floatingtrees Thanks, MTP has definitely been on our roadmap, and we’ve also been iterating on it continuously in our internal versions. Our current thinking is more along the lines of providing DeepSeek-V3 support in the spirit of #3392. Do you have any better ideas or suggestions?

tianyu-l

having some high level feedback before looking into details

tianyu-l · 2026-07-07T05:46:36Z

+        ({parallel_dims.tp}) and 2 * CP degree ({parallel_dims.cp}).
+        """
+
+    if parallelism.spmd_backend in ("full_dtensor", "spmd_types"):


We are in the migration to "spmd_types". Could you remove "default" and "full_dtensor" support from this new model?

Understood. So this new model only needs to support the spmd backend for now? Does this mean all torchtitan models will transition to supporting only the spmd backend in the future?

not addressed

There are still some issues when enabling --debug.spmd_typechecking. I am currently debugging this.

@tianyu-l To support spmd_types, we made the following changes:

Marked local topk_idxs / compress_topk_idxs as [B, S]-aligned metadata so they can be concatenated safely.

Disabled SPMD type checking around FlexAttention mask construction, where symbolic indexing is used.

Reworked compressor overlap logic to avoid {R} tensor in-place writes from {V} slices.

Replaced external Hadamard transform dependency with cached Hadamard matrix + F.linear.

Added sharding metadata for the hash-routing tid2eid buffer.

For details, please refer to:
85b3292

tianyu-l · 2026-07-07T05:48:04Z

+deepseek_v4_configs = {
+    "debugmodel": _debugmodel,
+}


why only debug model, can we include the full set?

I've included deepseek_v4_flash and deepseek_v4_pro now to expand the set. Just a heads-up: due to hardware constraints, I haven't actually tested them yet.

tianyu-l · 2026-07-07T05:50:13Z

    pipelining_fn: Callable | None
    post_optimizer_build_fn: Callable | None
    state_dict_adapter: type[BaseStateDictAdapter] | None
+    metrics_fn: Callable | None = None


@felipemello1 do you have a plan on in-model metric logging that we could share?

Let's add a TODO in dsv4 model code, and remove the model logging in this PR.

tianyu-l · 2026-07-07T05:56:16Z

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


I see similar contribution in #3864

Can we make sure the design in #3000 suits the need for both?

Yes, this implementation was referenced from #3864. I'll need to look closer into #3000 to see if the design can cover both use cases. Will update this thread once I have more clarity.

From a quick look, #3000 seems more like a MoE-specific interface, while here we would like to call the aux-loss capability through a common/shared interface that can also cover non-MoE losses such as the DSA indexer loss.

So for the current DeepSeek V4 use case, the #3864-style implementation seems more suitable to me.

I'll take a look at both

tianyu-l · 2026-07-07T05:58:04Z

could you please add verification (e.g. at least in PR summary) following https://github.com/pytorch/torchtitan/tree/main/scripts/checkpoint_conversion#comprehensive-check-kl-divergence

cc @shuhuayu

Thanks for the suggestion. Running the KL check with deepseek_v4_flash or deepseek_v4_pro as the baseline is not feasible for me due to hardware limits.

Would it be acceptable to run the same HF-vs-TorchTitan KL / max-diff comparison on a custom small DeepSeek V4 config instead?

Got it. @shuhuayu since you've done it for Kimi, could you help verify this PR gets to similar level of numerical parity?

Would it be acceptable to run the same HF-vs-TorchTitan KL / max-diff comparison on a custom small DeepSeek V4 config instead?

We also did a customized 16b kimi configs to compare numerics so it sounds good to me to do the same for dsv4. @tianyu-l, do you think we need to test it in the original 862b configs, doable but needs some additional setups.

Thanks. Could you share what config was used for the customized 16B Kimi numerical test? In particular, which dimensions were reduced, e.g. number of layers, hidden size, heads/head dim, MoE experts, or sequence length?
I can follow the same strategy for DSV4 and build a smaller custom config for the KL/numerical comparison.

I see. Small scale verification sounds fine.

Thanks. Could you share what config was used for the customized 16B Kimi numerical test?

I used the kimi-vl variant, see https://github.com/pytorch/torchtitan/pull/3532/changes#diff-5a13a40a31333dd8b095a114b15d9580c2744ac09211c6b551a02c2989032c9b

I used the kimi-vl variant, see https://github.com/pytorch/torchtitan/pull/3532/changes#diff-5a13a40a31333dd8b095a114b15d9580c2744ac09211c6b551a02c2989032c9b
Got it, will try it out soon. Thanks!

@tianyu-l I added a verification note based on the debug model comparison against the HuggingFace-style inference script.

The debug model now includes attention blocks with compress_ratios=(0, 0, 4, 128) and also covers the hash-routing MoE path, so it exercises the main algorithmic structures used by the current DeepSeek V4 implementation.

With the same weights and input tokens, the last-logits comparison between HF inference and TorchTitan gives cosine similarity:

cosine = 0.99999237

It is not bitwise identical mainly because:

The HF inference script uses a TileLang sparse attention kernel, while TorchTitan uses the current FlexAttention-based path, so minor rounding differences are expected.

The HF MoE path uses per-expert looped small matmuls, while TorchTitan uses grouped MM, which can also introduce small accumulation-order differences.

In addition, I also ran numerical validation across different parallelism strategies. With the same input, the loss comparison between FSDP and FSDP+EP+TP is:

[LOSS_COMPARE] Step-by-step loss comparison: [LOSS_COMPARE] Step Baseline Loss Test Loss Difference [LOSS_COMPARE] ---- ------------- --------- ---------- [LOSS_COMPARE] 1 8.082977294921875 8.082967758178711 -0.000010 [LOSS_COMPARE] 2 6.74751091003418 6.7473907470703125 -0.000120 [LOSS_COMPARE] 3 5.032241344451904 5.032076358795166 -0.000165 [LOSS_COMPARE] 4 4.599724769592285 4.599623680114746 -0.000101 [LOSS_COMPARE] 5 4.25457763671875 4.254598140716553 0.000021 [LOSS_COMPARE] 6 4.076618194580078 4.076696395874023 0.000078 [LOSS_COMPARE] 7 3.9123473167419434 3.9121904373168945 -0.000157 [LOSS_COMPARE] 8 3.7830865383148193 3.7829692363739014 -0.000117 [LOSS_COMPARE] 9 4.035467147827148 4.035507678985596 0.000041 [LOSS_COMPARE] 10 3.6527910232543945 3.6526951789855957 -0.000096 [LOSS_COMPARE] Summary statistics: [LOSS_COMPARE] Average baseline loss: 4.817734217643737 [LOSS_COMPARE] Average test loss: 4.81767156124115 [LOSS_COMPARE] Average difference: -0.000063

Given the debug model coverage, the high HF-vs-TorchTitan logits agreement, and the close distributed loss comparison, I think this provides a reasonable numerical verification for this PR.

sdmyzlp · 2026-07-31T03:37:20Z

+        with spmd.local():
+            n_local_groups = self.n_groups // (self.n_heads // o.shape[2])
+            o = o.view(bsz, seqlen, n_local_groups, -1)
+            _assert_spmd_attention_type(o, tp=spmd.S(2))
+            # wo_a is a Linear module; access its weight directly for the grouped
+            # einsum (not a standard Linear forward).
+            wo_a = self.wo_a.weight.view(n_local_groups, self.o_lora_rank, -1)
+            if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
+                spmd.assert_type(
+                    wo_a,
+                    {"dp": spmd.R, "cp": spmd.R, "tp": spmd.S(0)},
+                )
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)


How about makes this wo_a a BatchedLinear, for example or something similar:

class BatchedLinear(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): n_heads: int in_features: int out_features: int def __init__(self, config: Config): super().__init__() self.weight = nn.Parameter( torch.empty(config.n_heads, config.out_features, config.in_features) ) def forward(self, x: torch.Tensor) -> torch.Tensor: *prefix, H, D_in = x.shape x_h = x.reshape(-1, H, D_in).transpose(0, 1) # (H, T, D_in) out = torch.bmm(x_h, self.weight.transpose(-2, -1)) # (H, T, D_out) return out.transpose(0, 1).reshape(*prefix, H, -1)

this is cleaner but unclear to me how to represent per group attent ion output projection under spmd_types. looks to me it can invoke cross group communication?

sdmyzlp · 2026-07-31T08:24:31Z

+
+class DeepSeekV4Router(TokenChoiceTopKRouter):
+    @dataclass(kw_only=True, slots=True)
+    class Config(TokenChoiceTopKRouter.Config):


@tianyu-l Could we add a boolean hash field to TokenChioceTopKRouter.Config to remove the duplication below? Would that be a good way to go?

sdmyzlp · 2026-08-01T10:11:16Z

we have a mega PR here; meta-pytorch/attention-gym#225, we are going to break it up though and land in smaller chunks. It has attached some numbers

Hi @tianyu-l and @drisspg — Thanks for the pointer! We've closely studied the attention-gym CSA implementation and @floatingtrees' torchtitan override integration. Before we polish the DeepSeek-V4 attention further, we'd like to propose an adjustment to the inner_attention interface in this PR. We have a refactor on our side implementing it:

def forward(q, swa_k, cmp_k, idx_q, idx_k, idx_w, attn_sink, attention_masks):

swa_k is the sliding-window KV, cmp_k the compressed KV, and idx_* the raw indexer projections; the CSA top-k selection moves inside inner_attention. Several reasons for the modification:

Kernel-override friendliness. The current draft pre-concatenates swa_k + cmp_k before calling inner_attention, but fused kernels usually consume the sliding-window KV and the compressed KV as separated tensors, so an override would have to re-split the concatenated tensor just to call them. (one observation on the attention-gym side: the compressed_sparse_attention() entry point is a single monolithic call, which meant replacing the whole Attention module rather than inner_attention; exposing the modular pieces that already exist inside might make kernel integration smoother.)
Raw idx_* tensors seems to match fused loss kernel. Currently attention-gym's CSA has no backward for the indexer, so we can't tell for this part. Megatron-LM's fused indexer-loss implementation (which eventually delegates to cudnn.DSA) always recompute indexer scores; NPU kernels follow pretty much the same pattern.
Unified context-parallel handling. cmp_k and idx_k are K-like tensors and must be all-gathered under CP; listing them explicitly makes the all-gather a single, well-defined operation at the inner_attention boundary. (For comparison, on the current attention-gym API an integration would instead all-gather the uncompressed per-token tensors — this might also be a place where exposing a fine-grained interface would benefit, from our perspective.)
Removes the compressed-index bookkeeping entirely. With selection inside inner_attention, compressed indices never leave compressed-token space: the current draft's compress_topk_idxs + offset remapping into the concatenated KV space, the torch.cat([window_idxs, compress_idxs]) bookkeeping in Attention.forward, and the per-shard offset coupling that comes with a pre-concatenated KV all disappear.

Does this direction make sense? Thoughts or corrections would be much appreciated!

tianyu-l

I apologize for the delay in review.

The PR looks really good quality! Please see if my comments make sense. Thanks!

tianyu-l · 2026-08-14T06:02:41Z

+        key: torch.Tensor | None,
        rope_cache: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Apply a prepared RoPE cache to query and key.
+        *,
+        inverse: bool = False,


Do you think you could also help adapt the Helion kernel, as the "fast path"? https://github.com/pytorch/torchtitan/blob/main/torchtitan/overrides/helion_rope.py

Yes, I can take a look.

tianyu-l · 2026-08-14T06:05:00Z

+    idxs = torch.where(idxs >= causal_limit, -1, idxs + offset)
+    return idxs.unsqueeze(0).expand(bsz, -1, -1)
+
+class DSAIndexerAuxLoss(Module):


sorry for any back-and-forth, but let's maybe land #3864 first and rebase onto it

No worries at all. That sounds good to me. I’ll wait for #3864 to land, then rebase this PR on top of it and clean up any conflicts from there.

tianyu-l · 2026-08-14T06:12:37Z

+        return x
+
+
+class DeepSeekV4Model(Decoder):


Given that your MTP lands, why don't we integrate

Thanks. I added DeepSeek V4 MTP support in this PR, but it is not a direct drop-in of the DeepSeek V3 MTP path. DSV3 MTP runs on the regular single hidden stream, while DSV4 carries the HC hidden state as [B, L, hc_mult, D] through the main stack and each MTP depth needs to consume the previous HC state, roll the token input, run an inner DSV4 block, then collapse through its own HC head for prediction.

So the integration is a bit more model-specific than DSV3: the MTP block has to preserve the HC branches between depths, apply the valid mask after rolling packed positions, and avoid PP for now because the auxiliary outputs are not yet wired through the pipeline/chunked-loss path. I’m happy to integrate it here, but I kept the boundary explicit so it doesn’t look like the DSV3 MTP implementation can be reused unchanged.

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

+        return loss * self.coeff
+
+
+class DSAFlexAttention(FlexAttention):


Before we polish the DeepSeek-V4 attention further, we'd like to propose an adjustment to the inner_attention interface in this PR. We have a refactor on our side implementing it:
def forward(q, swa_k, cmp_k, idx_q, idx_k, idx_w, attn_sink, attention_masks):

This is certainly worth more discussion and I think it wouldn't block this PR. In particular, I worry about the composability with vLLM and our RL stack https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/models/attention.py

cc @drisspg how's attention-gym doing with sparse attention

@tianyu-l @pianpwk
One related question: the main reason I split the inner attention into CompressedSparseAttention, HeavilyCompressedAttention, and SlidingWindowAttention is that the sharding config / type-checking path does not seem to handle optional tensor arguments well. Conceptually, the implementation could be a single forward like forward(q, swa_k, cmp_k, idx_q, idx_k, idx_w, attn_sink, attention_masks), with some tensors being unused depending on compress_ratio, but those would need to be None for the non-compressed or heavily-compressed cases.

Do you know if there is a preferred way to represent optional tensor inputs in the sharding annotations today? If there is an existing pattern for this, I’m happy to consolidate the inner attention classes; otherwise keeping separate inner modules seems to make the sharding contract explicit for each path.

Thanks — understood that this doesn’t block the PR.

I looked at VLLMAttentionWrapper and vLLM’s DeepSeek-V4 implementation. The wrapper is designed for the standard (q, k, v) path used by GQAttention; vLLM itself handles DeepSeek-V4 through a dedicated DeepseekV4Attention, not through that generic seam.

So even without idx_*, DSV4’s inner_attention already has to move away from (q, k, v) toward something like (q, swa_k, cmp_k, indices, attention_sink). The only open question is whether we also pass idx_q/idx_k/idx_w.

If you have any thoughts on how much this would affect the RL/vLLM side, we’d be happy to hear them. Happy to adjust if there’s a cleaner way to keep the generic path untouched. @tianyu-l @drisspg

@sdmyzlp

If you have any thoughts on how much this would affect the RL/vLLM side, we’d be happy to hear them.

Thanks. Without looking into it too closely, my feeling is that we need to create dedicated vllm attention module, similar to what we do for Qwen 3.5's GDN https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/models/gdn.py

drisspg · 2026-08-14T09:11:57Z

@sdmyzlp sorry, just now seeing this. All of these points make sense, cc @floatingtrees for more opinions but I think we can do some refactoring to the attn gym interfaces

floatingtrees · 2026-08-17T23:48:51Z

@sdmyzlp Your points make a lot of sense. One thing that may help is that the version of CSA that landed on attention gym main is significantly stripped down. It takes in something like:
selected_attention(q, swa_k, cmp_k, indices, attention_sink)
where indices are the indices computed from idx_q, idx_k, idx_w, so it should be less clunky and easier to slot into inner_attention.

shuhuayu

Thanks for pr! left some comments.

shuhuayu · 2026-08-17T20:37:04Z

+    top_k = 6
+    n_hash_layers = 3
+    route_norm = True
+    route_scale = 1.5


should be 2.5 for pro model per hf config.json.

Thanks, fixed.

shuhuayu · 2026-08-17T22:36:19Z

+            skipped, in which case hidden states of shape ``[B, L, D]`` are
+            returned.
+        """
+        input_ids = tokens.detach().long()


not pp safe.

Thanks, fixed.

shuhuayu · 2026-08-17T22:36:31Z

+        """
+        input_ids = tokens.detach().long()
+        h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens
+        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)


not pp safe.

shuhuayu · 2026-08-17T22:36:59Z

+            layer = self.layers[str(i)]
+            h = layer(h, input_ids, attention_masks, positions)
+
+        h = self.hc_head(h)


not pp safe. maybe we should explicitly rejects pp if not supported.

shuhuayu · 2026-08-17T22:42:47Z

+
+        for i in range(self.n_main_layers):
+            layer = self.layers[str(i)]
+            h = layer(h, input_ids, attention_masks, positions)


input_ids is only necessary in the first n_hash_layers, should we carry it over all layers?

Good point. It is only used by the first n_hash_layers; later routers ignore it. I kept passing input_ids to all MoE layers to keep a uniform interface and sharding contract, since optional tensor args / None are not well supported by the current sharding config path.

shuhuayu · 2026-08-18T00:23:20Z

+            positions in ``[0, L)`` and padded entries are ``-1``.
+        """
+        window = min(seqlen, self.window_size)
+        q_idx = torch.arange(seqlen, device=device).unsqueeze(1)


raw indices are used, and positions, which contains sample boundary information, is only consumed by rope.

shuhuayu · 2026-08-18T03:22:09Z

+        with spmd.local():
+            n_local_groups = self.n_groups // (self.n_heads // o.shape[2])
+            o = o.view(bsz, seqlen, n_local_groups, -1)
+            _assert_spmd_attention_type(o, tp=spmd.S(2))
+            # wo_a is a Linear module; access its weight directly for the grouped
+            # einsum (not a standard Linear forward).
+            wo_a = self.wo_a.weight.view(n_local_groups, self.o_lora_rank, -1)
+            if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
+                spmd.assert_type(
+                    wo_a,
+                    {"dp": spmd.R, "cp": spmd.R, "tp": spmd.S(0)},
+                )
+        o = torch.einsum("bsgd,grd->bsgr", o, wo_a)


this is cleaner but unclear to me how to represent per group attent ion output projection under spmd_types. looks to me it can invoke cross group communication?

shuhuayu · 2026-08-18T08:25:54Z

+            tensor: Grouped tensor of shape ``[B, L // R, R, D]``.
+            value: Fill value for the first group's missing previous candidate.
+
+        Returns:
+            Tensor of shape ``[B, L // R, 2 * R, D]``.


the shapes should be from [B, L // R, R, D] to [B, L // R, 2*R, D//2]?

Yes, good catch. For the overlap path, the input to _overlap_transform has shape [B, L // R, R, 2 * head_dim], and the transform splits the last dim into previous/current candidates, then concatenates along the ratio dim, producing [B, L // R, 2 * R, head_dim]. So the current docstring is misleading; I’ll update it to make this explicit.

shuhuayu · 2026-08-18T08:44:59Z

+
+        row_max = comb.max(dim=-1, keepdim=True).values
+        comb = torch.exp(comb - row_max)
+        comb = comb / comb.sum(dim=-1, keepdim=True) + self.eps


why here eps is not in the denominator like in other cases?

Oh, good catch. Fixed.

shuhuayu · 2026-08-18T08:46:49Z

+        shape, dtype = x.size(), x.dtype
+        x = x.flatten(2).float()


we should do the same dtype convention in both hc_pre and hc_post?

Co-authored-by: Jonathan Zhou <96200235+floatingtrees@users.noreply.github.com>

sdmyzlp · 2026-08-20T14:10:33Z

It takes in something like: selected_attention(q, swa_k, cmp_k, indices, attention_sink) where indices are the indices computed from idx_q, idx_k, idx_w, so it should be less clunky and easier to slot into inner_attention.

@floatingtrees Yes, exactly, for a forward-only / indexer-loss-free path, that shape is clean and easy to slot into inner_attention.

One thing we’d like to ask you to evaluate is the indexer-loss training path. I looked through the attention-gym implementation and didn’t find indexer-loss handling there. If inner_attention only receives precomputed indices, we can do forward selection, but we cannot cleanly express the indexer-loss backward at the same boundary.

Megatron-LM’s cuDNN fused DSA (DeepSeek-V3.2) entry is:

fused_indexer_sparse_attn(
    query,      # main attention Q
    kv_full,    # main attention KV / compressed KV
    q_indexer,  # idx_q
    k_indexer,  # idx_k
    weights,    # idx_w
    ...
)

and is explicitly differentiable w.r.t. q_indexer, k_indexer, and weights; its backward returns grad_q_indexer, grad_k_indexer, and grad_weights. Megatron’s DSV4 prototype uses the same raw indexer inputs in FusedDSAIndexerLoss, even though that path is not yet cuDNN-fused.

To be transparent, one of our motivations for keeping idx_q/idx_k/idx_w in the inner_attention interface is also to make our fused-kernel integration easier. But I think the indexer-loss requirement is the more fundamental reason: the fused operator needs raw indexer projections in the same autograd scope as the main sparse attention.

Would you be open to evaluating the indexer-loss impact? If it holds, keeping idx_q/idx_k/idx_w in the inner_attention signature seems like the more future-proof contract.

floatingtrees · 2026-08-20T22:13:31Z

@sdmyzlp Ok, thanks for clarifying that. I agree that there doesn't seem to be a clean boundary for indexer loss under the current sparse attention API, but I'm not sure if it's general purpose enough to be placed in attention gym. Would it make sense if we returned the relevant tensors (such as query @ index_kv_full, where index_kv_full is kv_full after indexing), and leave loss function implementation to the caller? @drisspg do you think this would be an appropriate compromise for attention gym?

tianyu-l · 2026-08-21T08:27:46Z

+        scores_BLE: torch.Tensor,
+        expert_bias_E: torch.Tensor | None = None,
+        *,
+        input_ids: torch.Tensor | None = None,


I see that you are introducing router_kwargs, primarily for this input_ids field, but who's sending in this arg?

tianyu-l · 2026-08-21T08:29:04Z

+        return loss * self.coeff
+
+
+class DSAFlexAttention(FlexAttention):


@sdmyzlp

If you have any thoughts on how much this would affect the RL/vLLM side, we’d be happy to hear them.

Thanks. Without looking into it too closely, my feeling is that we need to create dedicated vllm attention module, similar to what we do for Qwen 3.5's GDN https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/models/gdn.py

Matrix-Z97 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners June 11, 2026 12:48

pytorch-bot Bot added the ciflow/8gpu label Jun 11, 2026

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

tianyu-l added this to 26H2 TorchTitan Development Jun 11, 2026

tianyu-l moved this to In Progress in 26H2 TorchTitan Development Jun 11, 2026

Matrix-Z97 closed this Jun 12, 2026

Matrix-Z97 force-pushed the deepseek_v4 branch from fda1fcc to 07828a4 Compare June 12, 2026 08:29

github-project-automation Bot moved this from In Progress to Done in 26H2 TorchTitan Development Jun 12, 2026

Matrix-Z97 reopened this Jun 12, 2026

Matrix-Z97 force-pushed the deepseek_v4 branch from daf4565 to 90c57d1 Compare June 12, 2026 08:39

tianyu-l moved this from Done to In Progress in 26H2 TorchTitan Development Jun 12, 2026

tianyu-l reviewed Jun 17, 2026

View reviewed changes

tianyu-l mentioned this pull request Jun 26, 2026

[WIP] DeepSeek-V3.2 model support #3800

Draft

1 task

floatingtrees reviewed Jun 26, 2026

View reviewed changes

Comment thread torchtitan/models/deepseek_v4/compressor.py Outdated

tianyu-l reviewed Jul 7, 2026

View reviewed changes

Matrix-Z97 requested review from SherlockNoMad, aditvenk, sanketpurandare and xmfan as code owners July 10, 2026 07:23

sdmyzlp reviewed Jul 31, 2026

View reviewed changes

Matrix-Z97 force-pushed the deepseek_v4 branch from 85b3292 to 7b76d5c Compare July 31, 2026 07:13

sdmyzlp reviewed Jul 31, 2026

View reviewed changes

tianyu-l requested a review from shuhuayu August 13, 2026 06:05

tianyu-l requested changes Aug 14, 2026

View reviewed changes

shuhuayu requested changes Aug 18, 2026

View reviewed changes

Matrix-Z97 force-pushed the deepseek_v4 branch from e198b27 to 390549e Compare August 20, 2026 07:38

Matrix-Z97 and others added 15 commits August 20, 2026 16:18

Add deepseek_v4 model support

2de803b

Update torchtitan/models/deepseek_v4/compressor.py

9df4af6

Co-authored-by: Jonathan Zhou <96200235+floatingtrees@users.noreply.github.com>

Refactor DeepSeek V4 sparse attention routing

1b9d756

Add DeepSeek V4 indexer aux loss

2d601cd

Update DeepSeek V4 configs and sharding

ed27afe

Refactor DeepSeek V4 attention, RoPE, and aux loss

bf569a7

Refactor DeepSeek V4 model implementation

fdda36c

Address DeepSeek V4 review comments

e23420d

Adapt DeepSeek V4 sparse paths for spmd_types

03733a1

Address DeepSeek V4 review feedback

0de3815

Remove common aux loss helper

7f0d84c

Add DeepSeek V4 sparse attention sharding support

a109b3c

Add DeepSeek V4 README

e6638dc

Add DeepSeek V4 MTP support

a9e4286

Add DeepSeek V4 MTP support

91d5e3d

Matrix-Z97 force-pushed the deepseek_v4 branch from 390549e to 91d5e3d Compare August 20, 2026 11:56

Address DeepSeek V4 review feedback

99222f6

tianyu-l reviewed Aug 21, 2026

View reviewed changes



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

		return loss * self.coeff


		class DSAFlexAttention(FlexAttention):

Uh oh!

Conversation

Matrix-Z97 commented Jun 11, 2026

Description

Main Changes

Parallelism Support

Additional Notes

Uh oh!

pytorch-bot Bot commented Jun 11, 2026

Uh oh!

Matrix-Z97 commented Jun 12, 2026

Uh oh!

tianyu-l commented Jun 12, 2026

Uh oh!

Matrix-Z97 commented Jun 12, 2026

Uh oh!

tianyu-l commented Jun 12, 2026

Uh oh!

Matrix-Z97 commented Jun 15, 2026

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!

floatingtrees commented Jun 29, 2026

Uh oh!

Matrix-Z97 commented Jul 2, 2026

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!

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!

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

Uh oh!

sdmyzlp Jul 31, 2026 •

edited

Loading

sdmyzlp Jul 31, 2026 •

edited

Loading

sdmyzlp commented Aug 1, 2026 •

edited

Loading