From d973fc19f766a4b65ed4c31eba5431fae720942d Mon Sep 17 00:00:00 2001 From: Angel Li Date: Thu, 27 Aug 2026 19:03:02 -0700 Subject: [PATCH 1/5] siwtch GDN to use attn-gym kernels --- .ci/docker/requirements.txt | 2 +- pyproject.toml | 2 +- tests/unit_tests/gpu/test_qwen3_5_deltanet.py | 114 +++++++++- torchtitan/experiments/rl/batch_invariance.py | 10 +- torchtitan/experiments/rl/models/gdn.py | 200 +++++++++--------- .../rl/tests/test_bitwise_parity.py | 9 +- torchtitan/models/kimi_k3/kda.py | 3 +- torchtitan/models/qwen3_5/gdn.py | 98 +++++---- 8 files changed, 284 insertions(+), 154 deletions(-) diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt index a2cec99127..7fb150c9ac 100644 --- a/.ci/docker/requirements.txt +++ b/.ci/docker/requirements.txt @@ -8,4 +8,4 @@ safetensors einops pillow spmd_types==0.2.5 -attn-gym[linear]==0.0.5 +attn-gym[linear]==0.0.6 diff --git a/pyproject.toml b/pyproject.toml index 8adfbd9fe2..b6b9bf0f2f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "einops", "pillow", "spmd_types==0.2.5", - "attn-gym[linear]==0.0.5", + "attn-gym[linear]==0.0.6", ] dynamic = ["version"] diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index 29a3cae33c..ac087c430b 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -9,7 +9,9 @@ import torch import torch.nn.functional as F +from attn_gym.linear import l2norm, recurrent_gdn from torch import nn + from torchtitan.models.common.attention import ( create_varlen_metadata_for_document, VarlenMetadata, @@ -103,8 +105,8 @@ def _reference_causal_conv1d_varlen( cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: - """Per-document depthwise causal conv + silu, matching the model's FLA - varlen conv (which is triton/CUDA-only). Patched over + """Per-document depthwise causal conv + silu, matching the model's Attention + Gym varlen conv (which is CUDA-only). Patched over ``gdn._causal_conv1d_varlen`` for CPU runs. """ conv_kernel_size = weight.shape[-1] @@ -401,7 +403,7 @@ def test_extracted_forward_matches_main(self): def _assert_packed_run_matches_per_document(self, model, x, positions, masks): """Packed forward under ``masks`` must equal stitched per-doc forwards. - The model's varlen conv is FLA (triton/CUDA-only); substitute the + The model's varlen conv is Attention Gym (CUDA-only); substitute the per-document torch reference for these CPU runs. The per-document forwards below take the non-varlen conv path, which runs on CPU. """ @@ -602,8 +604,112 @@ def test_fla_fused_recurrent_varlen_matches_independent_document_forwards(self): "fla_fused_recurrent", atol=2e-2, rtol=2e-2 ) + def test_batch_invariant_recurrent_matches_paged_attention_gym(self): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA is unavailable") + + from torchtitan.models.qwen3_5.gdn import _recurrent_gdn_fwd + + torch.manual_seed(42) + num_tokens, num_heads, key_dim, value_dim = 12, 4, 64, 64 + q = torch.randn( + 1, + num_tokens, + num_heads, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn_like(q) + v = torch.randn( + 1, + num_tokens, + num_heads, + value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + decay = -torch.rand( + 1, + num_tokens, + num_heads, + device="cuda", + dtype=torch.float32, + ) + update_gate = torch.rand( + 1, + num_tokens, + num_heads, + device="cuda", + dtype=torch.float32, + ) + cu_seqlens = torch.tensor([0, 5, 12], device="cuda", dtype=torch.int32) + + actual = _recurrent_gdn_fwd( + q, + k, + v, + decay, + update_gate, + cu_seqlens, + cu_seqlens.cpu(), + ) + + normalized_q = l2norm(q, cu_seqlens=cu_seqlens) + normalized_k = l2norm(k, cu_seqlens=cu_seqlens) + state_cache = torch.randn( + 5, + num_heads, + value_dim, + key_dim, + device="cuda", + dtype=torch.float32, + ) + prefix_end = 2 + prefix_cu_seqlens = torch.tensor( + [0, prefix_end], device="cuda", dtype=torch.int32 + ) + prefix_output, _ = recurrent_gdn( + normalized_q[:, :prefix_end], + normalized_k[:, :prefix_end], + v[:, :prefix_end], + decay[:, :prefix_end], + update_gate[:, :prefix_end], + state_cache, + cu_seqlens=prefix_cu_seqlens, + scale=key_dim**-0.5, + state_indices=torch.tensor([3], device="cuda", dtype=torch.int32), + has_initial_state=torch.tensor([False], device="cuda"), + ) + + state_indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + has_initial_state = torch.tensor([True, False], device="cuda") + remaining_cu_seqlens = torch.tensor( + [0, 5 - prefix_end, num_tokens - prefix_end], + device="cuda", + dtype=torch.int32, + ) + remaining_output, _ = recurrent_gdn( + normalized_q[:, prefix_end:], + normalized_k[:, prefix_end:], + v[:, prefix_end:], + decay[:, prefix_end:], + update_gate[:, prefix_end:], + state_cache, + cu_seqlens=remaining_cu_seqlens, + scale=key_dim**-0.5, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + expected = torch.cat( + (prefix_output, remaining_output), + dim=1, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + def test_varlen_offsets_are_fresh_per_deltanet_invocation(self): - """Successive DeltaNet invocations must not share FLA's cache key.""" + """Successive DeltaNet invocations must not share convolution metadata.""" torch.manual_seed(42) model = self._make_deltanet() x_TD = torch.randn(8, 4) diff --git a/torchtitan/experiments/rl/batch_invariance.py b/torchtitan/experiments/rl/batch_invariance.py index 3dd6f5bfb8..20ce18cda5 100644 --- a/torchtitan/experiments/rl/batch_invariance.py +++ b/torchtitan/experiments/rl/batch_invariance.py @@ -72,7 +72,15 @@ def patch_bmm_for_batch_invariance() -> None: global _batch_invariant_bmm_lib if _batch_invariant_bmm_lib is not None: return - from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant + try: + from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant + except ModuleNotFoundError as error: + if error.name not in { + "vllm.model_executor.determinism", + "vllm.model_executor.determinism.batch_invariant", + }: + raise + from vllm.model_executor.layers.batch_invariant import bmm_batch_invariant _batch_invariant_bmm_lib = torch.library.Library("aten", "IMPL") _batch_invariant_bmm_lib.impl("bmm", bmm_batch_invariant, "CUDA") diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py index 47348ba6d2..7f62e00ab5 100644 --- a/torchtitan/experiments/rl/models/gdn.py +++ b/torchtitan/experiments/rl/models/gdn.py @@ -6,44 +6,42 @@ """vLLM paged-cache adapter for TorchTitan's Gated DeltaNet. -The enclosing Qwen3.5 module owns all parameters. This adapter runs the same FLA -convolution and recurrence kernels as training while reading and updating vLLM's -paged convolution and SSM states. +The enclosing Qwen3.5 module owns all parameters. This adapter runs Attention +Gym's paging-aware convolution kernels. Batch-invariant recurrence and decode use +Attention Gym, while ordinary prefill uses FLA's parallel chunk kernel. Batch-invariant execution has two additional requirements: * The accumulated SSM cache state uses float32. Decode otherwise rounds the state through bfloat16 after every token, unlike a single prefill call. The convolution cache stays in model dtype because it only stores trailing input columns. -* Every recurrence receives a materialized initial state and ``cu_seqlens``. This - keeps FLA's Triton specialization identical for fresh prefill and resumed state. +* Batch-invariant recurrence uses the same Attention Gym scan as the trainer. -The recurrent kernel and materialized state are also used outside batch-invariant -mode because vLLM generation must support prefix continuation. SSM cache precision -and prefill kernel selection differ by mode. +Attention Gym updates the paged state pool directly. The non-batch-invariant FLA +prefill path retains its state gather and scatter. """ from dataclasses import dataclass import torch import torch.nn.functional as F - -# Using the trainer's FLA kernels is required for batch-invariant parity. -from fla.modules.convolution import ( - causal_conv1d as _fla_causal_conv1d, - causal_conv1d_update as _fla_causal_conv1d_update, +from attn_gym.linear import ( + causal_conv1d as _attn_gym_causal_conv1d, + causal_conv1d_decode as _attn_gym_causal_conv1d_decode, + l2norm as _attn_gym_l2norm, + recurrent_gdn as _attn_gym_recurrent_gdn, + recurrent_gdn_decode as _attn_gym_recurrent_gdn_decode, ) from fla.ops.gated_delta_rule import ( chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, - fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule, ) + from torchtitan.distributed.utils import is_in_batch_invariant_mode from torchtitan.protocols.module import Module # The recurrence mutates paged state and must run eager at a breakable cudagraph # split point. This decorator is inert when breakable capture is disabled. from vllm.compilation.breakable_cudagraph import eager_break_during_capture - from vllm.config import get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.model_executor.layers.mamba.abstract import MambaBase @@ -60,7 +58,7 @@ class VLLMInnerGatedDeltaNet(Module, MambaBase): """Paged-cache inner GDN implementation. The enclosing ``qwen3_5.gdn.GatedDeltaNet`` owns all parameters. This - module owns only vLLM cache plumbing and the FLA kernels. + module owns only vLLM cache plumbing and kernel dispatch. The enclosing module and vLLM cache are both head-sharded under tensor parallelism. Speculative decoding is not supported. @@ -86,6 +84,8 @@ def __init__(self, config: Config) -> None: self.num_speculative_tokens = ( speculative_config.num_speculative_tokens if speculative_config else 0 ) + if self.num_speculative_tokens != 0: + raise ValueError("Attention Gym GDN does not support speculative decoding.") self.num_k_heads = config.num_k_heads self.num_v_heads = config.num_v_heads @@ -108,11 +108,13 @@ def __init__(self, config: Config) -> None: self.local_num_v_heads = self.num_v_heads // self.tensor_parallel_size self.local_key_dim = self.local_num_k_heads * self.head_k_dim - if is_in_batch_invariant_mode(): - # The recurrent state accumulates in float32 inside FLA, so preserving - # it across scheduler calls avoids per-token bfloat16 rounding. The - # convolution cache only stores model-dtype input columns verbatim. - self.cache_config.mamba_ssm_cache_dtype = "float32" + if is_conv_state_dim_first(): + raise ValueError( + "Attention Gym GDN requires VLLM_SSM_CONV_STATE_LAYOUT=SD so " + "the paged convolution history has contiguous channels." + ) + + self.cache_config.mamba_ssm_cache_dtype = "float32" # vLLM populates this via the KV-cache allocator: (conv_state, ssm_state). self.kv_cache = (torch.tensor([]), torch.tensor([])) @@ -158,7 +160,7 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: def _split_qkv( self, mixed_qkv: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Split local fused channels and add FLA's singleton batch dim.""" + """Split local fused channels and add a singleton batch dim.""" num_tokens = mixed_qkv.shape[0] local_key_dim = self.local_key_dim query = ( @@ -185,17 +187,13 @@ def _run_recurrence( b: torch.Tensor, negative_exp_A: torch.Tensor, dt_bias: torch.Tensor, - initial_state: torch.Tensor, ssm_state: torch.Tensor, slot_indices: torch.Tensor, cu_seqlens: torch.Tensor, - use_chunk: bool, + has_initial_state: torch.Tensor | None, + batch_invariant: bool, ) -> torch.Tensor: - """Run FLA's gated-delta rule and update the selected paged SSM slots. - - ``use_chunk`` selects the parallel chunk kernel over the sequential - recurrent kernel; the caller sets it for non-batch-invariant prefill. - """ + """Run the selected recurrence and update the paged SSM slots.""" query, key, value = self._split_qkv(conv_output) # Grouped-value heads: expand q/k to match the value head count. if query.shape[2] != value.shape[2]: @@ -206,17 +204,41 @@ def _run_recurrence( decay = (negative_exp_A * F.softplus(a.float() + dt_bias)).unsqueeze(0) update_gate = torch.sigmoid(b).unsqueeze(0) - # Non-batch-invariant prefill/mixed batches use the parallel chunk kernel - # (far faster than the O(seqlen) sequential recurrence). Batch-invariant - # mode keeps the recurrent kernel so the generator matches the trainer's - # recurrent forward bitwise; pure decode is a single token, where chunking - # gains nothing and the recurrent kernel stays cudagraph-capturable. - gated_delta_rule = ( - _fla_chunk_gated_delta_rule - if use_chunk - else _fla_fused_recurrent_gated_delta_rule + if batch_invariant: + query = _attn_gym_l2norm(query, cu_seqlens=cu_seqlens) + key = _attn_gym_l2norm(key, cu_seqlens=cu_seqlens) + output, _ = _attn_gym_recurrent_gdn( + query, + key, + value, + decay, + update_gate, + ssm_state, + cu_seqlens=cu_seqlens, + scale=self.head_k_dim**-0.5, + state_indices=slot_indices, + has_initial_state=has_initial_state, + # Triton autotuning breaks batch invariance. + autotune=False, + ) + return output + + num_sequences = slot_indices.numel() + initial_state = conv_output.new_zeros( + num_sequences, + self.local_num_v_heads, + self.head_k_dim, + self.head_v_dim, + dtype=torch.float32, ) - output, final_state = gated_delta_rule( + if has_initial_state is None: + initial_state.copy_(ssm_state[slot_indices].transpose(-1, -2)) + else: + resumed_slots = slot_indices[has_initial_state] + initial_state[has_initial_state] = ssm_state[resumed_slots].transpose( + -1, -2 + ) + output, final_state = _fla_chunk_gated_delta_rule( query, key, value, @@ -227,6 +249,7 @@ def _run_recurrence( cu_seqlens=cu_seqlens, use_qk_l2norm_in_kernel=True, ) + assert final_state is not None ssm_state[slot_indices] = final_state.transpose(-1, -2).to(ssm_state.dtype) return output @@ -269,13 +292,8 @@ def _forward( state_indices = gdn_metadata.non_spec_state_indices_tensor assert state_indices is not None ssm_state = self.kv_cache[1] - # vLLM's default conv-state layout keeps channels innermost (.., W-1, C); - # FLA needs channels first (.., C, W-1), so transpose that case. - conv_state = ( - self.kv_cache[0] - if is_conv_state_dim_first() - else self.kv_cache[0].transpose(-1, -2) - ) + conv_state = self.kv_cache[0] + assert conv_bias is None negative_exp_A = -torch.exp(A_log.float()) dt_bias = dt_bias.float() num_decodes = gdn_metadata.num_decodes @@ -293,21 +311,13 @@ def _forward( decode_slots = state_indices[:num_decodes] if num_decodes > 0: # pure decode, or mixed prefill decode - decode_conv_state = conv_state[decode_slots] - # vLLM stores the trailing W - 1 inputs. FLA expects W entries but - # ignores the first, so prepend a zero column and drop it on write. - zero_padding = decode_conv_state.new_zeros( - decode_conv_state.shape[0], decode_conv_state.shape[1], 1 - ) - conv_cache = torch.cat([zero_padding, decode_conv_state], dim=-1) - decode_conv_output, conv_cache = _fla_causal_conv1d_update( - mixed_qkv[:num_decode_tokens], - conv_cache, - weight=conv_weight, - bias=conv_bias, + decode_conv_output = _attn_gym_causal_conv1d_decode( + mixed_qkv[:num_decode_tokens].contiguous(), + conv_weight, + conv_state, activation="silu", + state_indices=decode_slots, ) - conv_state[decode_slots] = conv_cache[..., 1:].to(conv_state.dtype) conv_output[:num_decode_tokens] = decode_conv_output prefill_slots = None @@ -316,6 +326,7 @@ def _forward( assert gdn_metadata.prefill_state_indices is not None prefill_slots = gdn_metadata.prefill_state_indices prefill_has_initial_state = gdn_metadata.prefill_has_initial_state + assert prefill_has_initial_state is not None prefill_start = num_decode_tokens if num_decodes > 0 else 0 # cu_seqlens must be 0-based within the prefill slice that the conv # kernel receives (mixed_qkv[prefill_start:]). @@ -340,85 +351,74 @@ def _forward( prefill_has_initial_state.any() ) conv_initial_state = mixed_qkv.new_zeros( - num_prefill_sequences, mixed_qkv.shape[1], self.conv_kernel_size + num_prefill_sequences, + self.conv_kernel_size - 1, + mixed_qkv.shape[1], ) # Fresh prefills keep zero state; prefix-cache continuations restore # only the sequence slots identified by vLLM metadata. if has_continuations: resumed_slots = prefill_slots[prefill_has_initial_state] - conv_initial_state[prefill_has_initial_state, :, 1:] = conv_state[ + conv_initial_state[prefill_has_initial_state] = conv_state[ resumed_slots ] - prefill_conv_output, conv_final_state = _fla_causal_conv1d( + prefill_conv_output, conv_final_state = _attn_gym_causal_conv1d( mixed_qkv[prefill_start:num_actual_tokens].unsqueeze(0), - weight=conv_weight, - bias=conv_bias, + conv_weight, activation="silu", cu_seqlens=prefill_cu_seqlens, initial_state=conv_initial_state, - output_final_state=True, + return_final_state=True, ) - conv_state[prefill_slots] = conv_final_state[..., 1:].to(conv_state.dtype) + conv_state[prefill_slots] = conv_final_state.to(conv_state.dtype) conv_output[prefill_start:num_actual_tokens] = prefill_conv_output.squeeze( 0 ) - # Recurrence over the whole batch in one call. Decode (T=1) and prefill - # (T>1) sequences run together, delimited by cu_seqlens; each gets a - # materialized fp32 initial state. The kernel processes each sequence - # independently from its cu_seqlens entry and initial-state row, so a - # single call covers the whole batch (chunk or recurrent, chosen below). + # Recurrence over the whole batch in one call. The batch-invariant path + # addresses the paged SSM pool directly; ordinary prefill uses FLA. cu_seqlens = gdn_metadata.non_spec_query_start_loc[: num_sequences + 1] if num_prefills == 0: - # Pure decode is captured in a CUDA graph, which forbids host syncs and - # data-dependent (boolean-mask) indexing. Every decode sequence resumes, - # so gather its paged SSM state directly with the integer slot indices. all_slots = decode_slots - initial_state = ( - ssm_state[all_slots].transpose(-1, -2).to(torch.float32).contiguous() - ) + has_initial_state = None else: - # Prefill / mixed runs eager at the graph break, so boolean-mask gather - # and the host sync it implies are allowed. A sequence resumes from - # paged state iff it is a decode (always) or a prefill prefix-cache - # continuation; fresh prefills keep a zero initial state. all_slots = ( torch.cat([decode_slots, prefill_slots]) if num_decodes > 0 else prefill_slots ) - resumes_from_cache = torch.zeros( + has_initial_state = torch.ones( num_sequences, dtype=torch.bool, device=mixed_qkv.device ) - resumes_from_cache[:num_decodes] = True if prefill_has_initial_state is not None: - resumes_from_cache[num_decodes:] = prefill_has_initial_state - initial_state = conv_output.new_zeros( - num_sequences, - self.local_num_v_heads, - self.head_k_dim, - self.head_v_dim, - dtype=torch.float32, - ) - initial_state[resumes_from_cache] = ( - ssm_state[all_slots[resumes_from_cache]] - .transpose(-1, -2) - .to(torch.float32) + has_initial_state[num_decodes:] = prefill_has_initial_state + + batch_invariant = is_in_batch_invariant_mode() + if num_prefills == 0 and not batch_invariant: + _attn_gym_recurrent_gdn_decode( + conv_output[:num_decode_tokens], + a[:num_decode_tokens].unsqueeze(0), + b[:num_decode_tokens].unsqueeze(0), + A_log.float(), + dt_bias, + ssm_state, + all_slots, + scale=self.head_k_dim**-0.5, + out=output[:num_decode_tokens].unsqueeze(0), ) - # Non-batch-invariant prefill/mixed batches take the parallel chunk kernel; - # batch-invariant mode and pure decode take the recurrent kernel. - use_chunk = num_prefills > 0 and not is_in_batch_invariant_mode() + return + recurrent_output = self._run_recurrence( conv_output, a, b, negative_exp_A, dt_bias, - initial_state, ssm_state, all_slots, cu_seqlens, - use_chunk, + has_initial_state, + batch_invariant=batch_invariant, ) output[:num_actual_tokens] = recurrent_output[0, :num_actual_tokens].to( output.dtype diff --git a/torchtitan/experiments/rl/tests/test_bitwise_parity.py b/torchtitan/experiments/rl/tests/test_bitwise_parity.py index abca28be0a..65cd8728ba 100644 --- a/torchtitan/experiments/rl/tests/test_bitwise_parity.py +++ b/torchtitan/experiments/rl/tests/test_bitwise_parity.py @@ -633,16 +633,17 @@ def setUpClass(cls): if hf_path: config.hf_assets_path = hf_path - from torchtitan.tools.utils import has_cuda_capability + from torchtitan.tools.utils import get_cuda_flash_attention_impl - if has_cuda_capability(9, 0): + flash_attention_impl = get_cuda_flash_attention_impl() + if flash_attention_impl is not None: from torch.nn.attention import ( activate_flash_attention_impl, current_flash_attention_impl, ) - if current_flash_attention_impl() != "FA3": - activate_flash_attention_impl("FA3") + if current_flash_attention_impl() != flash_attention_impl: + activate_flash_attention_impl(flash_attention_impl) # Enable batch-invariant mode BEFORE init_distributed set_batch_invariance(config.trainer.debug.batch_invariant) diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index ebc7618bcc..552d7929e4 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from attn_gym.linear.kda import bound_gate, chunk_kda from attn_gym.linear.kda.fwd.triton.l2norm_fwd import l2norm -from attn_gym.linear.kda.short_conv import causal_conv1d +from attn_gym.linear.short_conv import causal_conv1d from torch import nn from torchtitan.models.common.attention import AttentionMasksType, VarlenMetadata @@ -96,7 +96,6 @@ def forward( A_log_H.float(), dt_bias_HK.float(), lower_bound=self.lower_bound, - impl="fused", ) output_1THV, _ = chunk_kda( l2norm(q_1THK), diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py index aa774351dd..7624bf25d3 100644 --- a/torchtitan/models/qwen3_5/gdn.py +++ b/torchtitan/models/qwen3_5/gdn.py @@ -6,13 +6,20 @@ """Gated DeltaNet modules for Qwen3.5.""" +# Tensor dimensions: B = batch, T = tokens, N = heads, K = key dimension, +# V = value dimension, S = state slots, D = channels. + from dataclasses import dataclass from typing import Literal import spmd_types as spmd import torch import torch.nn.functional as F -from fla.modules.conv.triton.ops import CausalConv1dFunction +from attn_gym.linear import ( + causal_conv1d as _attn_gym_causal_conv1d, + l2norm as _attn_gym_l2norm, + recurrent_gdn as _attn_gym_recurrent_gdn, +) from fla.ops.gated_delta_rule import ( chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule, @@ -35,7 +42,6 @@ spmd.register_local_autograd_function(ChunkGatedDeltaRuleFunction) spmd.register_local_autograd_function(FusedRecurrentFunction) -spmd.register_local_autograd_function(CausalConv1dFunction) @spmd.local_map( @@ -53,28 +59,19 @@ def _causal_conv1d_varlen( cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor | None, ) -> torch.Tensor: - """FLA depthwise causal conv with per-document resets (CUDA-only). + """Depthwise causal conv with per-document resets (CUDA-only). A pure-torch per-document reference lives in ``tests/unit_tests/gpu/test_qwen3_5_deltanet.py``. """ - if cu_seqlens_cpu is None: - raise ValueError( - "Qwen3.5 FLA varlen conv requires a CPU cu_seqlens tensor. " - "Build VarlenMetadata with include_host_offsets=True." - ) - - from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d - - out_1TD, _ = _fla_causal_conv1d( - x=x_TD.unsqueeze(0), - weight=weight.squeeze(1), - bias=None, + del cu_seqlens_cpu + out_1TD = _attn_gym_causal_conv1d( + x_TD.unsqueeze(0), + weight.squeeze(1), activation="silu", - backend="triton", cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, ) + assert isinstance(out_1TD, torch.Tensor) return out_1TD.squeeze(0) @@ -109,9 +106,9 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: "torchtitan::recurrent_gdn_fwd", mutates_args=(), device_types="cuda" ) def _recurrent_gdn_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + q_BTNK: torch.Tensor, + k_BTNK: torch.Tensor, + v_BTNV: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor, @@ -119,43 +116,62 @@ def _recurrent_gdn_fwd( ) -> torch.Tensor: """Run the batch-invariant GDN recurrent forward kernel. - The vLLM generator must use the recurrent kernel for per-token decode. The - trainer uses the same kernel with a materialized float32 initial state and - varlen metadata so its forward is bitwise identical to generation. + The vLLM generator uses Attention Gym's paging-aware recurrent kernel for + per-token decode. The trainer uses the same recurrence with a materialized + float32 initial state and varlen metadata so its forward is bitwise identical + to generation. """ num_sequences = int(cu_seqlens.numel()) - 1 - initial_state = q.new_zeros( - num_sequences, - q.shape[2], - q.shape[3], - v.shape[3], + # state_cache_SNVK: [num_sequences + 1, N, V, K]. + state_cache_SNVK = q_BTNK.new_empty( + num_sequences + 1, + q_BTNK.shape[2], + v_BTNV.shape[3], + q_BTNK.shape[3], dtype=torch.float32, ) - output, _ = _fla_fused_recurrent_gated_delta_rule( - q, - k, - v, + state_indices = torch.arange( + 1, + num_sequences + 1, + dtype=torch.int32, + device=q_BTNK.device, + ) + has_initial_state = torch.zeros( + num_sequences, + dtype=torch.bool, + device=q_BTNK.device, + ) + # FLA normalizes Q/K inside its kernels. Attention Gym's recurrent GDN + # expects normalized inputs, so apply the same normalization explicitly. + normalized_q_BTNK = _attn_gym_l2norm(q_BTNK, cu_seqlens=cu_seqlens) + normalized_k_BTNK = _attn_gym_l2norm(k_BTNK, cu_seqlens=cu_seqlens) + out_BTNV, _ = _attn_gym_recurrent_gdn( + normalized_q_BTNK, + normalized_k_BTNK, + v_BTNV, g, - beta=beta, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, + beta, + state_cache_SNVK, cu_seqlens=cu_seqlens, + scale=q_BTNK.shape[-1] ** -0.5, + state_indices=state_indices, + has_initial_state=has_initial_state, + autotune=False, ) - return output.to(q.dtype) + return out_BTNV.to(q_BTNK.dtype) @_recurrent_gdn_fwd.register_fake def _recurrent_gdn_fwd_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, + q_BTNK: torch.Tensor, + k_BTNK: torch.Tensor, + v_BTNV: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: - return torch.empty_like(v, dtype=q.dtype) + return torch.empty_like(v_BTNV, dtype=q_BTNK.dtype) @torch.library.custom_op( From 6f1e074ab6304e3ba5630cfdc32932dba0fb4490 Mon Sep 17 00:00:00 2001 From: drisspg Date: Mon, 31 Aug 2026 21:42:59 -0700 Subject: [PATCH 2/5] Replace Qwen3.5 FLA paths with Attention Gym Use Attention Gym for Qwen3.5 GDN training, batch-invariant backward recomputation, recurrent execution, decode, and paged chunk prefill. The paged prefill path advances vLLM SSM cache slots in place, eliminating the previous gather/chunk/scatter copies and preserving the shared [slots, H, V, K] layout across prefill and decode. Remove the now-unused FLA backend selection and dependency, normalize Q/K explicitly at the Attention Gym boundary, and update debug configurations to the fused backend K=V=128 contract while retaining two key heads for TP=2 coverage. On GB200, isolated paged-prefill measurements reduced GPU time by 18.2-31.7% across five T=1024-8192, N=1-64 cases. A matched end-to-end vLLM debug-model comparison at prompt=128, generation=64, and batch sizes 1-16 improved output throughput by 1.5-6.4%. ```bash PYTHONPATH=/home/drisspg/meta/attention-gym-paged-gdn:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/attention-gym-paged-gdn:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/cpu/test_state_dict_keys.py tests/unit_tests/cpu/test_integration_test_definitions.py tests/unit_tests/cpu/test_train_spec.py -x PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym-paged-gdn torchrun --nproc-per-node=2 -m pytest torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen -v -s ``` --- .ci/docker/requirements-vlm.txt | 1 - .../workflows/integration_test_8gpu_rl.yaml | 8 +- pyproject.toml | 2 +- scripts/ci/pytorch_ci_test_runner.sh | 4 +- tests/unit_tests/gpu/test_qwen3_5_deltanet.py | 242 +++++++--------- .../test_qwen3_5_mrope_positions.py | 22 +- torchtitan/experiments/rl/models/gdn.py | 75 ++--- .../experiments/rl/models/vllm_wrapper.py | 2 +- torchtitan/models/kimi_k3/README.md | 2 +- torchtitan/models/qwen3_5/README.md | 2 +- torchtitan/models/qwen3_5/__init__.py | 25 +- torchtitan/models/qwen3_5/gdn.py | 270 ++++++------------ torchtitan/models/qwen3_5/sharding.py | 4 +- torchtitan_recipes/tests/__init__.py | 2 +- torchtitan_recipes/tests/models.py | 3 - 15 files changed, 235 insertions(+), 429 deletions(-) diff --git a/.ci/docker/requirements-vlm.txt b/.ci/docker/requirements-vlm.txt index 9b1e557669..e82b2e33ba 100644 --- a/.ci/docker/requirements-vlm.txt +++ b/.ci/docker/requirements-vlm.txt @@ -2,4 +2,3 @@ av einops pillow torchvision -flash-linear-attention diff --git a/.github/workflows/integration_test_8gpu_rl.yaml b/.github/workflows/integration_test_8gpu_rl.yaml index 73968a31d1..f950b08d48 100644 --- a/.github/workflows/integration_test_8gpu_rl.yaml +++ b/.github/workflows/integration_test_8gpu_rl.yaml @@ -80,14 +80,10 @@ jobs: --extra-index-url https://download.pytorch.org/whl/nightly/cu130 \ --index-strategy unsafe-best-match - # 4. Install the pinned GDN kernels without replacing nightly PyTorch. - uv pip install --no-deps \ - "git+https://github.com/fla-org/flash-linear-attention.git@v0.5.2" - - # 5. Make the checkout importable for subprocesses spawned by the test. + # 4. Make the checkout importable for subprocesses spawned by the test. export PYTHONPATH="$PWD:${PYTHONPATH:-}" - # 6. Download HF model checkpoint for tests + # 5. Download HF model checkpoint for tests MODEL_PATH=$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('Qwen/Qwen3-0.6B'))") sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded" diff --git a/pyproject.toml b/pyproject.toml index b6b9bf0f2f..41b673d48e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,5 +77,5 @@ markers = [ [tool.pyrefly] python-version = "3.11" project-excludes = ["torchtitan/experiments", "**/tests/**"] -replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "fla.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"] # optional dependencies +replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"] # optional dependencies search-path = ["../pytorch"] # local built pytorch diff --git a/scripts/ci/pytorch_ci_test_runner.sh b/scripts/ci/pytorch_ci_test_runner.sh index 2bfa6ffb69..8957c2620b 100755 --- a/scripts/ci/pytorch_ci_test_runner.sh +++ b/scripts/ci/pytorch_ci_test_runner.sh @@ -41,9 +41,7 @@ case "$COMMAND" in model_tests) # qwen3_5_fsdp+tp+varlen_attn+per_op_sac: varlen attention needs # flash_attn_interface/FA3, which the PyTorch CI image does not install - # and which is unavailable on its A10G (sm86) runners anyway. Excluded - # here rather than disabled in models.py so torchtitan's own CI, whose - # image ships flash-linear-attention, keeps running it. + # and which is unavailable on its A10G (sm86) runners anyway. python -m tests.integration_tests.run_tests \ --test_suite models \ --exclude "qwen3_5_moe_fsdp+tp+ep+pp,qwen3_5_fsdp+tp+varlen_attn+per_op_sac" \ diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index ac087c430b..b020a32a7d 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -17,12 +17,12 @@ VarlenMetadata, ) -# Tensor shape suffixes: B batch, L seq len, H heads, K query/key head dim, +# Tensor shape suffixes: B batch, L seq len, H heads, K key head dim, # V value head dim. def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor: - """L2 norm using rsqrt(sum(x^2) + eps), not x/max(norm, eps) like F.normalize, to match FLA kernel.""" + """Match Attention Gym's rsqrt-based L2 normalization.""" return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) @@ -36,7 +36,7 @@ def _torch_native_gated_delta( """Standalone math reference for the gated delta rule recurrence. Sequential O(seqlen) loop -- far too slow for training; kept here as the - numerical baseline for the FLA kernels. + numerical baseline for the fused kernels. Args: q_BLHK, k_BLHK: (batch, seq, num_heads, key_head_dim) @@ -45,7 +45,7 @@ def _torch_native_gated_delta( beta_BLH: (batch, seq, num_heads) -- update gate in (0, 1) Returns: - output: (batch, seq, n_heads, value_head_dim) + output: (batch, seq, num_heads, value_head_dim) """ B, L, H, K = q_BLHK.shape V = v_BLHV.shape[-1] @@ -103,28 +103,27 @@ def _reference_causal_conv1d_varlen( x_TD: torch.Tensor, weight: torch.Tensor, cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: """Per-document depthwise causal conv + silu, matching the model's Attention Gym varlen conv (which is CUDA-only). Patched over ``gdn._causal_conv1d_varlen`` for CPU runs. """ conv_kernel_size = weight.shape[-1] - out_segments_1TD: list[torch.Tensor] = [] - cu_seqlens_list = cu_seqlens_cpu.tolist() + out_segments_BTD: list[torch.Tensor] = [] + cu_seqlens_list = cu_seqlens.tolist() for start, end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:], strict=False): - x_segment_1DT = F.pad( + x_segment_BDT = F.pad( x_TD[start:end].transpose(0, 1).unsqueeze(0), [conv_kernel_size - 1, 0], ) - out_segment_1TD = F.conv1d( - x_segment_1DT, + out_segment_BTD = F.conv1d( + x_segment_BDT, weight, None, groups=weight.size(0), ).transpose(1, 2) - out_segments_1TD.append(out_segment_1TD) - return F.silu(torch.cat(out_segments_1TD, dim=1)).squeeze(0) + out_segments_BTD.append(out_segment_BTD) + return F.silu(torch.cat(out_segments_BTD, dim=1)).squeeze(0) class ReferenceGatedDeltaKernel(nn.Module): @@ -145,7 +144,6 @@ def forward( beta_TH: torch.Tensor, *, cu_seqlens: torch.Tensor | None = None, - cu_seqlens_cpu: torch.Tensor | None = None, ) -> torch.Tensor: if xq_THK.shape[1] != xv_THV.shape[1]: assert xv_THV.shape[1] % xq_THK.shape[1] == 0 @@ -153,19 +151,18 @@ def forward( xq_THK = xq_THK.repeat_interleave(repeat, dim=1) xk_THK = xk_THK.repeat_interleave(repeat, dim=1) - xq_1THK = xq_THK.unsqueeze(0) - xk_1THK = xk_THK.unsqueeze(0) - xv_1THV = xv_THV.unsqueeze(0) - g_1TH = g_TH.unsqueeze(0) - beta_1TH = beta_TH.unsqueeze(0) + xq_BLHK = xq_THK.unsqueeze(0) + xk_BLHK = xk_THK.unsqueeze(0) + xv_BLHV = xv_THV.unsqueeze(0) + g_BLH = g_TH.unsqueeze(0) + beta_BLH = beta_TH.unsqueeze(0) if cu_seqlens is None: return _torch_native_gated_delta( - xq_1THK, xk_1THK, xv_1THV, g_1TH, beta_1TH + xq_BLHK, xk_BLHK, xv_BLHV, g_BLH, beta_BLH ).squeeze(0) - assert cu_seqlens_cpu is not None return _torch_native_gated_delta_varlen( - xq_1THK, xk_1THK, xv_1THV, g_1TH, beta_1TH, cu_seqlens_cpu + xq_BLHK, xk_BLHK, xv_BLHV, g_BLH, beta_BLH, cu_seqlens.cpu() ).squeeze(0) @@ -180,8 +177,7 @@ def test_flex_masks_ignore_padding_position_resets(self): ) from exc with torch.device("meta"): - build_config, max_context_length = qwen3_5_configs["debugmodel"] - model = build_config("flex", seq_len=max_context_length).build() + model = qwen3_5_configs["debugmodel"]("flex").build() positions = torch.tensor([0, 1, 2, 0, 0], dtype=torch.int32) with mock.patch.object(Decoder, "get_attention_masks", return_value=None): @@ -199,8 +195,7 @@ def test_flex_masks_include_delta_net_varlen_metadata(self): ) from exc with torch.device("meta"): - build_config, max_context_length = qwen3_5_configs["debugmodel"] - model = build_config("flex", seq_len=max_context_length).build() + model = qwen3_5_configs["debugmodel"]("flex").build() positions = torch.tensor([0, 1, 0, 1, 2], dtype=torch.int32) full_attention_mask = mock.sentinel.full_attention_mask @@ -221,10 +216,7 @@ def test_flex_masks_include_delta_net_varlen_metadata(self): def _make_deltanet( self, *, - # None builds the model with the default FLA kernel config, then swaps - # in ReferenceGatedDeltaKernel so the model runs on CPU without FLA - # triton kernels. - backend: str | None = None, + use_fused: bool = False, dim: int = 4, key_head_dim: int = 2, value_head_dim: int = 2, @@ -280,11 +272,7 @@ def conv(channels: int) -> Conv1d.Config: conv_k=conv(key_dim), conv_v=conv(value_dim), inner_gated_delta_net=InnerGatedDeltaNet.Config( - kernel=( - GatedDeltaKernel.Config() - if backend is None - else GatedDeltaKernel.Config(backend=backend) - ), + kernel=GatedDeltaKernel.Config(), ), norm=RMSNormGated.Config(dim=value_head_dim), out_proj=Linear.Config( @@ -293,7 +281,7 @@ def conv(channels: int) -> Conv1d.Config: bias=False, ), ).build() - if backend is None: + if not use_fused: model.inner_gated_delta_net.kernel = ReferenceGatedDeltaKernel() model = model.to(device=device, dtype=dtype) @@ -316,14 +304,8 @@ def _main_forward_reference(self, model, x_TD, attention_masks=None): """Run the current main-branch GatedDeltaNet forward structure.""" num_tokens = x_TD.shape[0] cu_seqlens = None - cu_seqlens_cpu = None if attention_masks is not None: cu_seqlens = attention_masks.cu_seq_q.clone() - cu_seqlens_cpu = torch.tensor( - attention_masks.cu_seq_q_host, - dtype=cu_seqlens.dtype, - device="cpu", - ) def causal_conv(tensor, conv): if cu_seqlens is not None: @@ -331,7 +313,6 @@ def causal_conv(tensor, conv): tensor, conv.weight, cu_seqlens, - cu_seqlens_cpu, ) tensor = F.pad( tensor.transpose(0, 1).unsqueeze(0), @@ -373,7 +354,6 @@ def causal_conv(tensor, conv): decay_TH, update_gate_TH, cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, ) output_THV = model.norm(output_THV, gate_THV) return model.out_proj(output_THV.reshape(num_tokens, -1)) @@ -432,7 +412,7 @@ def test_varlen_matches_independent_document_forwards(self): attention_masks = create_varlen_metadata_for_document( positions, - include_host_offsets=True, + include_host_offsets=False, ) self._assert_packed_run_matches_per_document( model, x_TD, positions, attention_masks @@ -446,9 +426,6 @@ def test_get_attention_masks_pairs_flex_mask_with_deltanet_offsets(self): """ from torch.nn.attention.flex_attention import BlockMask - # torchtitan.models.qwen3_5 imports the FLA (flash-linear-attention) - # kernels at module scope. FLA is a triton/CUDA-only optional - # dependency, so skip instead of erroring on environments without it. try: from torchtitan.models.qwen3_5 import model_registry except ModuleNotFoundError as exc: @@ -507,8 +484,8 @@ def test_get_attention_masks_pairs_flex_mask_with_deltanet_offsets(self): (0, 3, 5, 10), ) - def _assert_fla_varlen_matches_per_document( - self, backend: str, *, atol: float, rtol: float + def _assert_fused_varlen_matches_per_document( + self, *, atol: float, rtol: float ) -> None: if not torch.cuda.is_available(): raise unittest.SkipTest("CUDA is unavailable") @@ -516,16 +493,15 @@ def _assert_fla_varlen_matches_per_document( device = "cuda" dtype = torch.bfloat16 torch.manual_seed(42) - # Mirror the debug model's GatedDeltaNet dims so the FLA Triton kernels - # accept the shapes; n_value_heads > n_key_heads also exercises the - # grouped-query head expansion inside the kernel. + # The fused chunk kernel requires 128-wide heads. Unequal key and value + # head counts also exercise grouped-head execution. model = self._make_deltanet( - backend=backend, + use_fused=True, dim=256, - key_head_dim=64, - value_head_dim=64, - num_key_heads=2, - num_value_heads=4, + key_head_dim=128, + value_head_dim=128, + num_key_heads=1, + num_value_heads=2, conv_kernel_size=4, device=device, dtype=dtype, @@ -563,7 +539,13 @@ def _assert_fla_varlen_matches_per_document( dtype=torch.int32, device=device, ) - x_TD = torch.randn(positions.shape[0], 256, device=device, dtype=dtype) + x_TD = torch.randn( + positions.shape[0], + 256, + device=device, + dtype=dtype, + requires_grad=True, + ) attention_masks = create_varlen_metadata_for_document( positions, @@ -571,9 +553,8 @@ def _assert_fla_varlen_matches_per_document( ) actual = model(x_TD, attention_masks) - # Reference: run each document on its own (non-varlen path) and stitch - # the outputs back. Matching this proves the FLA varlen kernels reset - # recurrent state at document boundaries instead of bleeding across them. + # Reference: run each document on its own and stitch the outputs back. + # Matching proves packed execution resets state at document boundaries. expected = torch.empty_like(actual) doc_starts = (positions == 0).nonzero(as_tuple=True)[0].tolist() ends = doc_starts[1:] + [positions.shape[0]] @@ -584,25 +565,23 @@ def _assert_fla_varlen_matches_per_document( self.assertTrue( torch.allclose(actual, expected, rtol=rtol, atol=atol), msg=( - f"{backend}: varlen output diverged from per-document forwards " + "varlen output diverged from per-document forwards " f"(max abs diff {max_diff:.3e}, atol {atol}, rtol {rtol}). " "Cross-document state bleed produces diffs on the order of the " "output magnitude, far larger than bf16 kernel noise." ), ) - def test_fla_chunked_varlen_matches_independent_document_forwards(self): - # bf16 tolerance absorbs the differing chunk boundaries between the - # packed varlen run and the per-document runs; tighten once confirmed on - # GPU (the failure message reports the observed max diff). - self._assert_fla_varlen_matches_per_document( - "fla_chunked", atol=2e-2, rtol=2e-2 - ) + actual.float().square().mean().backward() + self.assertIsNotNone(x_TD.grad) + self.assertTrue(torch.isfinite(x_TD.grad).all()) + for parameter in model.parameters(): + self.assertIsNotNone(parameter.grad) + self.assertTrue(torch.isfinite(parameter.grad).all()) - def test_fla_fused_recurrent_varlen_matches_independent_document_forwards(self): - self._assert_fla_varlen_matches_per_document( - "fla_fused_recurrent", atol=2e-2, rtol=2e-2 - ) + def test_fused_varlen_matches_independent_document_forwards(self): + # BF16 tolerance absorbs differing packed and per-document chunk boundaries. + self._assert_fused_varlen_matches_per_document(atol=2e-2, rtol=2e-2) def test_batch_invariant_recurrent_matches_paged_attention_gym(self): if not torch.cuda.is_available(): @@ -611,11 +590,17 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): from torchtitan.models.qwen3_5.gdn import _recurrent_gdn_fwd torch.manual_seed(42) - num_tokens, num_heads, key_dim, value_dim = 12, 4, 64, 64 + num_tokens, num_key_heads, num_value_heads, key_dim, value_dim = ( + 12, + 1, + 2, + 128, + 128, + ) q = torch.randn( 1, num_tokens, - num_heads, + num_key_heads, key_dim, device="cuda", dtype=torch.bfloat16, @@ -624,7 +609,7 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): v = torch.randn( 1, num_tokens, - num_heads, + num_value_heads, value_dim, device="cuda", dtype=torch.bfloat16, @@ -632,34 +617,35 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): decay = -torch.rand( 1, num_tokens, - num_heads, + num_value_heads, device="cuda", dtype=torch.float32, ) update_gate = torch.rand( 1, num_tokens, - num_heads, + num_value_heads, device="cuda", dtype=torch.float32, ) + for tensor in (q, k, v, decay, update_gate): + tensor.requires_grad_() cu_seqlens = torch.tensor([0, 5, 12], device="cuda", dtype=torch.int32) - actual = _recurrent_gdn_fwd( + actual = torch.compile(_recurrent_gdn_fwd, fullgraph=True)( q, k, v, decay, update_gate, cu_seqlens, - cu_seqlens.cpu(), ) normalized_q = l2norm(q, cu_seqlens=cu_seqlens) normalized_k = l2norm(k, cu_seqlens=cu_seqlens) state_cache = torch.randn( 5, - num_heads, + num_value_heads, value_dim, key_dim, device="cuda", @@ -669,18 +655,19 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): prefix_cu_seqlens = torch.tensor( [0, prefix_end], device="cuda", dtype=torch.int32 ) - prefix_output, _ = recurrent_gdn( - normalized_q[:, :prefix_end], - normalized_k[:, :prefix_end], - v[:, :prefix_end], - decay[:, :prefix_end], - update_gate[:, :prefix_end], - state_cache, - cu_seqlens=prefix_cu_seqlens, - scale=key_dim**-0.5, - state_indices=torch.tensor([3], device="cuda", dtype=torch.int32), - has_initial_state=torch.tensor([False], device="cuda"), - ) + with torch.no_grad(): + prefix_output, _ = recurrent_gdn( + normalized_q[:, :prefix_end], + normalized_k[:, :prefix_end], + v[:, :prefix_end], + decay[:, :prefix_end], + update_gate[:, :prefix_end], + state_cache, + cu_seqlens=prefix_cu_seqlens, + scale=key_dim**-0.5, + state_indices=torch.tensor([3], device="cuda", dtype=torch.int32), + has_initial_state=torch.tensor([False], device="cuda"), + ) state_indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) has_initial_state = torch.tensor([True, False], device="cuda") @@ -689,18 +676,19 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): device="cuda", dtype=torch.int32, ) - remaining_output, _ = recurrent_gdn( - normalized_q[:, prefix_end:], - normalized_k[:, prefix_end:], - v[:, prefix_end:], - decay[:, prefix_end:], - update_gate[:, prefix_end:], - state_cache, - cu_seqlens=remaining_cu_seqlens, - scale=key_dim**-0.5, - state_indices=state_indices, - has_initial_state=has_initial_state, - ) + with torch.no_grad(): + remaining_output, _ = recurrent_gdn( + normalized_q[:, prefix_end:], + normalized_k[:, prefix_end:], + v[:, prefix_end:], + decay[:, prefix_end:], + update_gate[:, prefix_end:], + state_cache, + cu_seqlens=remaining_cu_seqlens, + scale=key_dim**-0.5, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) expected = torch.cat( (prefix_output, remaining_output), dim=1, @@ -708,46 +696,10 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): torch.testing.assert_close(actual, expected, rtol=0, atol=0) - def test_varlen_offsets_are_fresh_per_deltanet_invocation(self): - """Successive DeltaNet invocations must not share convolution metadata.""" - torch.manual_seed(42) - model = self._make_deltanet() - x_TD = torch.randn(8, 4) - positions = torch.tensor( - [0, 1, 2, 0, 1, 2, 3, 4], - dtype=torch.int32, - ) - attention_masks = create_varlen_metadata_for_document( - positions, - include_host_offsets=True, - ) - captured_cu_seqlens = [] - - def record_cu_seqlens(x_TD, weight, cu_seqlens, cu_seqlens_cpu): - captured_cu_seqlens.append(cu_seqlens) - return _reference_causal_conv1d_varlen( - x_TD, - weight, - cu_seqlens, - cu_seqlens_cpu, - ) - - with mock.patch( - "torchtitan.models.qwen3_5.gdn._causal_conv1d_varlen", - side_effect=record_cu_seqlens, - ): - model(x_TD, attention_masks) - model(x_TD, attention_masks) - - # Main runs separate Q/K/V convolutions, so each invocation uses the - # same cloned offsets three times. - self.assertEqual(len(captured_cu_seqlens), 6) - first_invocation = captured_cu_seqlens[0] - second_invocation = captured_cu_seqlens[3] - self.assertTrue(all(x is first_invocation for x in captured_cu_seqlens[:3])) - self.assertTrue(all(x is second_invocation for x in captured_cu_seqlens[3:])) - self.assertIsNot(first_invocation, attention_masks.cu_seq_q) - self.assertIsNot(second_invocation, first_invocation) + actual.float().square().mean().backward() + for tensor in (q, k, v, decay, update_gate): + self.assertIsNotNone(tensor.grad) + self.assertTrue(torch.isfinite(tensor.grad).all()) if __name__ == "__main__": diff --git a/tests/unit_tests/test_qwen3_5_mrope_positions.py b/tests/unit_tests/test_qwen3_5_mrope_positions.py index ee59c5c1fd..4ab1c8a988 100644 --- a/tests/unit_tests/test_qwen3_5_mrope_positions.py +++ b/tests/unit_tests/test_qwen3_5_mrope_positions.py @@ -18,6 +18,8 @@ resolution lives in ``forward`` or in ``preprocess_inputs``. """ +import subprocess +import sys import unittest import torch @@ -25,9 +27,6 @@ def _build_config_modules(): - # torchtitan.models.qwen3_5 imports the FLA (flash-linear-attention) - # kernels at module scope. FLA is a triton/CUDA-only optional dependency, - # so skip instead of erroring on environments without it. try: from torchtitan.config import ParallelismConfig from torchtitan.distributed.parallel_dims import ParallelDims @@ -43,8 +42,8 @@ class _RecordingLayer(nn.Module): """Layer stub that records the positions it is handed and passes x through. The mrope/positions resolution is independent of the layer internals, so - stubbing the layers keeps these tests on CPU without the FLA kernels while - still exercising the real ``preprocess_inputs`` and ``forward`` glue. + stubbing the layers keeps these tests on CPU while still exercising the real + ``preprocess_inputs`` and ``forward`` glue. """ def __init__(self, sink: dict): @@ -57,6 +56,19 @@ def forward(self, x, attention_masks=None, positions=None): class TestQwen35MRoPEPositions(unittest.TestCase): + def test_model_import_does_not_require_fla(self): + script = ( + "import builtins\n" + "original_import = builtins.__import__\n" + "def without_fla(name, globals=None, locals=None, fromlist=(), level=0):\n" + " if level == 0 and (name == 'fla' or name.startswith('fla.')):\n" + " raise ModuleNotFoundError('blocked fla import')\n" + " return original_import(name, globals, locals, fromlist, level)\n" + "builtins.__import__ = without_fla\n" + "import torchtitan.models.qwen3_5\n" + ) + subprocess.run([sys.executable, "-c", script], check=True) + def _build_stub_model(self): model_registry, ParallelDims, ParallelismConfig = _build_config_modules() # varlen backend keeps mask construction to pure tensor ops (no flex diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py index 7f62e00ab5..4dc2c737d4 100644 --- a/torchtitan/experiments/rl/models/gdn.py +++ b/torchtitan/experiments/rl/models/gdn.py @@ -7,8 +7,7 @@ """vLLM paged-cache adapter for TorchTitan's Gated DeltaNet. The enclosing Qwen3.5 module owns all parameters. This adapter runs Attention -Gym's paging-aware convolution kernels. Batch-invariant recurrence and decode use -Attention Gym, while ordinary prefill uses FLA's parallel chunk kernel. +Gym's paging-aware convolution and GDN kernels. Batch-invariant execution has two additional requirements: @@ -17,8 +16,9 @@ cache stays in model dtype because it only stores trailing input columns. * Batch-invariant recurrence uses the same Attention Gym scan as the trainer. -Attention Gym updates the paged state pool directly. The non-batch-invariant FLA -prefill path retains its state gather and scatter. +Decode, recurrent execution, and chunked prefill all update the paged SSM state +pool directly. Convolution prefill still materializes its much smaller ``W - 1`` +history until Attention Gym exposes a paged varlen convolution operation. """ from dataclasses import dataclass @@ -26,14 +26,12 @@ import torch import torch.nn.functional as F from attn_gym.linear import ( - causal_conv1d as _attn_gym_causal_conv1d, - causal_conv1d_decode as _attn_gym_causal_conv1d_decode, - l2norm as _attn_gym_l2norm, - recurrent_gdn as _attn_gym_recurrent_gdn, - recurrent_gdn_decode as _attn_gym_recurrent_gdn_decode, -) -from fla.ops.gated_delta_rule import ( - chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, + causal_conv1d, + causal_conv1d_decode, + l2norm, + paged_chunk_gdn, + recurrent_gdn, + recurrent_gdn_decode, ) from torchtitan.distributed.utils import is_in_batch_invariant_mode @@ -195,19 +193,13 @@ def _run_recurrence( ) -> torch.Tensor: """Run the selected recurrence and update the paged SSM slots.""" query, key, value = self._split_qkv(conv_output) - # Grouped-value heads: expand q/k to match the value head count. - if query.shape[2] != value.shape[2]: - num_repeats = value.shape[2] // query.shape[2] - query = query.repeat_interleave(num_repeats, dim=2) - key = key.repeat_interleave(num_repeats, dim=2) - decay = (negative_exp_A * F.softplus(a.float() + dt_bias)).unsqueeze(0) update_gate = torch.sigmoid(b).unsqueeze(0) + query = l2norm(query, cu_seqlens=cu_seqlens) + key = l2norm(key, cu_seqlens=cu_seqlens) if batch_invariant: - query = _attn_gym_l2norm(query, cu_seqlens=cu_seqlens) - key = _attn_gym_l2norm(key, cu_seqlens=cu_seqlens) - output, _ = _attn_gym_recurrent_gdn( + output, _ = recurrent_gdn( query, key, value, @@ -223,35 +215,18 @@ def _run_recurrence( ) return output - num_sequences = slot_indices.numel() - initial_state = conv_output.new_zeros( - num_sequences, - self.local_num_v_heads, - self.head_k_dim, - self.head_v_dim, - dtype=torch.float32, - ) - if has_initial_state is None: - initial_state.copy_(ssm_state[slot_indices].transpose(-1, -2)) - else: - resumed_slots = slot_indices[has_initial_state] - initial_state[has_initial_state] = ssm_state[resumed_slots].transpose( - -1, -2 - ) - output, final_state = _fla_chunk_gated_delta_rule( + return paged_chunk_gdn( query, key, value, decay, - beta=update_gate, - initial_state=initial_state, - output_final_state=True, + update_gate, + ssm_state, + slot_indices, cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True, + has_initial_state=has_initial_state, + scale=self.head_k_dim**-0.5, ) - assert final_state is not None - ssm_state[slot_indices] = final_state.transpose(-1, -2).to(ssm_state.dtype) - return output # The decorator makes this an eager graph-split point during breakable capture. # The caller-owned output has a stable address across graph replays. @@ -311,7 +286,7 @@ def _forward( decode_slots = state_indices[:num_decodes] if num_decodes > 0: # pure decode, or mixed prefill decode - decode_conv_output = _attn_gym_causal_conv1d_decode( + decode_conv_output = causal_conv1d_decode( mixed_qkv[:num_decode_tokens].contiguous(), conv_weight, conv_state, @@ -362,7 +337,7 @@ def _forward( conv_initial_state[prefill_has_initial_state] = conv_state[ resumed_slots ] - prefill_conv_output, conv_final_state = _attn_gym_causal_conv1d( + prefill_conv_output, conv_final_state = causal_conv1d( mixed_qkv[prefill_start:num_actual_tokens].unsqueeze(0), conv_weight, activation="silu", @@ -376,7 +351,7 @@ def _forward( ) # Recurrence over the whole batch in one call. The batch-invariant path - # addresses the paged SSM pool directly; ordinary prefill uses FLA. + # addresses the paged SSM pool directly. cu_seqlens = gdn_metadata.non_spec_query_start_loc[: num_sequences + 1] if num_prefills == 0: all_slots = decode_slots @@ -395,7 +370,7 @@ def _forward( batch_invariant = is_in_batch_invariant_mode() if num_prefills == 0 and not batch_invariant: - _attn_gym_recurrent_gdn_decode( + recurrent_gdn_decode( conv_output[:num_decode_tokens], a[:num_decode_tokens].unsqueeze(0), b[:num_decode_tokens].unsqueeze(0), @@ -440,11 +415,9 @@ def forward( *, key_head_dim: int, value_head_dim: int, - cu_seqlens_host: tuple[int, ...] | None = None, + use_packed_sequence: bool = False, ) -> torch.Tensor: """Run the flattened vLLM cache operation on rank-local tensors.""" - del cu_seqlens, cu_seqlens_host - assert key_head_dim == self.head_k_dim assert value_head_dim == self.head_v_dim mixed_qkv_TC = torch.cat([query_TC, key_TC, value_TC], dim=-1) diff --git a/torchtitan/experiments/rl/models/vllm_wrapper.py b/torchtitan/experiments/rl/models/vllm_wrapper.py index 134aa9d423..fe401b2ffb 100644 --- a/torchtitan/experiments/rl/models/vllm_wrapper.py +++ b/torchtitan/experiments/rl/models/vllm_wrapper.py @@ -59,7 +59,7 @@ def _replace_vllm_layer_configs(model_config): # These modules inspect the breakable-cudagraph environment at import time. # Defer imports until vLLM constructs the model, after the generator has set # that environment. Import the GDN adapter only for hybrid models so other - # models do not acquire FLA as an optional dependency. + # models do not acquire its vLLM-specific dependencies. from torchtitan.experiments.rl.models.attention import VLLMAttentionWrapper new_layers = [] diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 2a0f09205e..765d72825d 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -8,7 +8,7 @@ Attention (MLA) decoder with LatentMoE and a MoonViT-V2 vision encoder. Install the additional dependencies: ```bash -pip install av einops pillow torchvision flash-linear-attention +pip install av einops pillow torchvision ``` ## Architecture diff --git a/torchtitan/models/qwen3_5/README.md b/torchtitan/models/qwen3_5/README.md index 0ab86d5f3c..2147416772 100644 --- a/torchtitan/models/qwen3_5/README.md +++ b/torchtitan/models/qwen3_5/README.md @@ -24,7 +24,7 @@ Note: the diagram shows each patch mapping to one vision token. In practice, the Install the additional dependencies: ```bash -pip install av torchvision flash-linear-attention +pip install av torchvision ``` ## Model Variants diff --git a/torchtitan/models/qwen3_5/__init__.py b/torchtitan/models/qwen3_5/__init__.py index 96205d3b8a..34bffdd89e 100644 --- a/torchtitan/models/qwen3_5/__init__.py +++ b/torchtitan/models/qwen3_5/__init__.py @@ -38,13 +38,7 @@ from torchtitan.protocols.model_spec import ModelSpec -from .gdn import ( - GatedDeltaBackend, - GatedDeltaKernel, - GatedDeltaNet, - InnerGatedDeltaNet, - RMSNormGated, -) +from .gdn import GatedDeltaKernel, GatedDeltaNet, InnerGatedDeltaNet, RMSNormGated from .model import OffsetRMSNorm, Qwen35Attention, Qwen35Model, Qwen35TransformerBlock from .parallelize import parallelize_qwen3_5 @@ -265,7 +259,6 @@ def _qwen35_deltanet_config( value_head_dim: int, layer_id: int, conv_kernel_size: int = 4, - fla_backend: GatedDeltaBackend = "fla_chunked", ) -> GatedDeltaNet.Config: """Build a fully-specified GatedDeltaNet.Config.""" key_dim = n_key_heads * key_head_dim @@ -302,7 +295,7 @@ def _conv(channels: int) -> Conv1d.Config: conv_k=_conv(key_dim), conv_v=_conv(value_dim), inner_gated_delta_net=InnerGatedDeltaNet.Config( - kernel=GatedDeltaKernel.Config(backend=fla_backend), + kernel=GatedDeltaKernel.Config(), ), norm=RMSNormGated.Config( dim=value_head_dim, @@ -333,7 +326,6 @@ def _build_qwen35_layers( value_head_dim: int, full_attention_interval: int = 4, attn_backend: str, - fla_backend: GatedDeltaBackend = "fla_chunked", ) -> list[Qwen35TransformerBlock.Config]: """Build per-layer configs for dense Qwen3.5 models.""" layers = [] @@ -362,7 +354,6 @@ def _build_qwen35_layers( key_head_dim=key_head_dim, value_head_dim=value_head_dim, layer_id=layer_id, - fla_backend=fla_backend, ) if not is_full else None @@ -404,7 +395,6 @@ def _build_qwen35_moe_layers( value_head_dim: int, full_attention_interval: int = 4, attn_backend: str, - fla_backend: GatedDeltaBackend = "fla_chunked", moe_comm_backend: str = "standard", non_blocking_capacity_factor: float | None = None, ) -> list[Qwen35TransformerBlock.Config]: @@ -435,7 +425,6 @@ def _build_qwen35_moe_layers( key_head_dim=key_head_dim, value_head_dim=value_head_dim, layer_id=layer_id, - fla_backend=fla_backend, ) if not is_full else None @@ -518,9 +507,8 @@ def _debugmodel(attn_backend: str, *, seq_len: int) -> Qwen35Model.Config: hidden_dim=512, n_key_heads=2, n_value_heads=4, - key_head_dim=64, - value_head_dim=64, - fla_backend="fla_chunked", + key_head_dim=128, + value_head_dim=128, ), vision_encoder=_qwen35_vision_encoder_config( dim=256, @@ -583,10 +571,9 @@ def _debugmodel_moe( shared_expert_hidden_dim=256, n_key_heads=2, n_value_heads=4, - key_head_dim=64, - value_head_dim=64, + key_head_dim=128, + value_head_dim=128, moe_comm_backend=moe_comm_backend, - fla_backend="fla_chunked", ), vision_encoder=_qwen35_vision_encoder_config( dim=256, diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py index 7624bf25d3..daba81b9ba 100644 --- a/torchtitan/models/qwen3_5/gdn.py +++ b/torchtitan/models/qwen3_5/gdn.py @@ -6,26 +6,17 @@ """Gated DeltaNet modules for Qwen3.5.""" -# Tensor dimensions: B = batch, T = tokens, N = heads, K = key dimension, -# V = value dimension, S = state slots, D = channels. +# Shape suffixes: +# T = packed tokens, D = model dimension, C = projection channels, +# H = attention heads, K = query/key head dimension, V = value head dimension, +# S = state slots, W = convolution kernel width. from dataclasses import dataclass -from typing import Literal import spmd_types as spmd import torch import torch.nn.functional as F -from attn_gym.linear import ( - causal_conv1d as _attn_gym_causal_conv1d, - l2norm as _attn_gym_l2norm, - recurrent_gdn as _attn_gym_recurrent_gdn, -) -from fla.ops.gated_delta_rule import ( - chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, - fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule, -) -from fla.ops.gated_delta_rule.chunk import ChunkGatedDeltaRuleFunction -from fla.ops.gated_delta_rule.fused_recurrent import FusedRecurrentFunction +from attn_gym.linear import causal_conv1d, chunk_gdn, l2norm, recurrent_gdn from torch import nn from torchtitan.distributed.utils import is_in_batch_invariant_mode @@ -33,23 +24,12 @@ from torchtitan.models.common.attention import VarlenMetadata from torchtitan.protocols.module import Module -# Shape suffixes: -# T = packed tokens, D = model dimension, C = projection channels, -# H = attention heads, K = query/key head dimension, V = value head dimension, -# W = convolution kernel width. - -GatedDeltaBackend = Literal["fla_chunked", "fla_fused_recurrent"] - -spmd.register_local_autograd_function(ChunkGatedDeltaRuleFunction) -spmd.register_local_autograd_function(FusedRecurrentFunction) - @spmd.local_map( in_types=( {"dp": spmd.S(0), "tp": spmd.S(1)}, {"dp": spmd.R, "tp": spmd.S(0)}, {"dp": spmd.V, "tp": spmd.R}, - {"dp": spmd.V, "tp": spmd.R}, ), out_types={"dp": spmd.S(0), "tp": spmd.S(1)}, ) @@ -57,22 +37,20 @@ def _causal_conv1d_varlen( x_TD: torch.Tensor, weight: torch.Tensor, cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor | None, ) -> torch.Tensor: """Depthwise causal conv with per-document resets (CUDA-only). A pure-torch per-document reference lives in ``tests/unit_tests/gpu/test_qwen3_5_deltanet.py``. """ - del cu_seqlens_cpu - out_1TD = _attn_gym_causal_conv1d( + out_BTD = causal_conv1d( x_TD.unsqueeze(0), weight.squeeze(1), activation="silu", cu_seqlens=cu_seqlens, ) - assert isinstance(out_1TD, torch.Tensor) - return out_1TD.squeeze(0) + assert isinstance(out_BTD, torch.Tensor) + return out_BTD.squeeze(0) class RMSNormGated(Module): @@ -106,13 +84,12 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: "torchtitan::recurrent_gdn_fwd", mutates_args=(), device_types="cuda" ) def _recurrent_gdn_fwd( - q_BTNK: torch.Tensor, - k_BTNK: torch.Tensor, - v_BTNV: torch.Tensor, + q_BTHK: torch.Tensor, + k_BTHK: torch.Tensor, + v_BTHV: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: """Run the batch-invariant GDN recurrent forward kernel. @@ -122,62 +99,57 @@ def _recurrent_gdn_fwd( to generation. """ num_sequences = int(cu_seqlens.numel()) - 1 - # state_cache_SNVK: [num_sequences + 1, N, V, K]. - state_cache_SNVK = q_BTNK.new_empty( + # state_cache_SHVK: [num_sequences + 1, H, V, K]. + state_cache_SHVK = q_BTHK.new_empty( num_sequences + 1, - q_BTNK.shape[2], - v_BTNV.shape[3], - q_BTNK.shape[3], + v_BTHV.shape[2], + v_BTHV.shape[3], + q_BTHK.shape[3], dtype=torch.float32, ) state_indices = torch.arange( 1, num_sequences + 1, dtype=torch.int32, - device=q_BTNK.device, + device=q_BTHK.device, ) has_initial_state = torch.zeros( num_sequences, dtype=torch.bool, - device=q_BTNK.device, + device=q_BTHK.device, ) - # FLA normalizes Q/K inside its kernels. Attention Gym's recurrent GDN - # expects normalized inputs, so apply the same normalization explicitly. - normalized_q_BTNK = _attn_gym_l2norm(q_BTNK, cu_seqlens=cu_seqlens) - normalized_k_BTNK = _attn_gym_l2norm(k_BTNK, cu_seqlens=cu_seqlens) - out_BTNV, _ = _attn_gym_recurrent_gdn( - normalized_q_BTNK, - normalized_k_BTNK, - v_BTNV, + # The recurrent operator consumes normalized Q/K. + normalized_q_BTHK = l2norm(q_BTHK, cu_seqlens=cu_seqlens) + normalized_k_BTHK = l2norm(k_BTHK, cu_seqlens=cu_seqlens) + out_BTHV, _ = recurrent_gdn( + normalized_q_BTHK, + normalized_k_BTHK, + v_BTHV, g, beta, - state_cache_SNVK, + state_cache_SHVK, cu_seqlens=cu_seqlens, - scale=q_BTNK.shape[-1] ** -0.5, + scale=q_BTHK.shape[-1] ** -0.5, state_indices=state_indices, has_initial_state=has_initial_state, autotune=False, ) - return out_BTNV.to(q_BTNK.dtype) + return out_BTHV.to(q_BTHK.dtype) @_recurrent_gdn_fwd.register_fake def _recurrent_gdn_fwd_fake( - q_BTNK: torch.Tensor, - k_BTNK: torch.Tensor, - v_BTNV: torch.Tensor, + q_BTHK: torch.Tensor, + k_BTHK: torch.Tensor, + v_BTHV: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: - return torch.empty_like(v_BTNV, dtype=q_BTNK.dtype) + return torch.empty_like(v_BTHV, dtype=q_BTHK.dtype) -@torch.library.custom_op( - "torchtitan::chunk_gdn_bwd", mutates_args=(), device_types="cuda" -) -def _chunk_gdn_bwd( +def _chunk_gdn_gradients( grad_output: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -185,56 +157,37 @@ def _chunk_gdn_bwd( g: torch.Tensor, beta: torch.Tensor, cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Recompute the parallel GDN chunk kernel and return its gradients.""" with torch.enable_grad(): inputs = tuple( tensor.detach().requires_grad_(True) for tensor in (q, k, v, g, beta) ) - output = _fla_chunk_gated_delta_rule( - inputs[0], - inputs[1], + normalized_q = l2norm(inputs[0], cu_seqlens=cu_seqlens) + normalized_k = l2norm(inputs[1], cu_seqlens=cu_seqlens) + output, _ = chunk_gdn( + normalized_q, + normalized_k, inputs[2], inputs[3], inputs[4], - use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - )[0] + scale=inputs[0].shape[-1] ** -0.5, + impl="fused", + ) grad_q, grad_k, grad_v, grad_g, grad_beta = torch.autograd.grad( output, inputs, grad_output ) return grad_q, grad_k, grad_v, grad_g, grad_beta -@_chunk_gdn_bwd.register_fake -def _chunk_gdn_bwd_fake( - grad_output: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - cu_seqlens: torch.Tensor, - cu_seqlens_cpu: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - return ( - torch.empty_like(q), - torch.empty_like(k), - torch.empty_like(v), - torch.empty_like(g), - torch.empty_like(beta), - ) - - def _recurrent_gdn_setup_context(ctx, inputs, output) -> None: ctx.save_for_backward(*inputs) def _recurrent_gdn_backward(ctx, grad_output): - q, k, v, g, beta, cu_seqlens, cu_seqlens_cpu = ctx.saved_tensors - grads = _chunk_gdn_bwd( + q, k, v, g, beta, cu_seqlens = ctx.saved_tensors + grads = _chunk_gdn_gradients( grad_output, q, k, @@ -242,9 +195,8 @@ def _recurrent_gdn_backward(ctx, grad_output): g, beta, cu_seqlens, - cu_seqlens_cpu, ) - return (*grads, None, None) + return (*grads, None) _recurrent_gdn_fwd.register_autograd( @@ -253,25 +205,19 @@ def _recurrent_gdn_backward(ctx, grad_output): class GatedDeltaKernel(Module): - """Stateless dispatch to the configured FLA gated delta kernel. - - Provides a module boundary for the sharding code to wrap forward with - DTensor-to-local conversion -- same pattern as FlexAttention. Handles Q/K - head expansion for grouped linear attention internally so that - repeat_interleave runs on local tensors under TP. A pure-torch reference - implementation lives in ``tests/unit_tests/gpu/test_qwen3_5_deltanet.py``; - it is far too slow for training use. + """Run GDN on rank-local tensors. + + This module provides the boundary that sharding wraps with DTensor-to-local + conversion. A pure-torch reference implementation lives in + ``tests/unit_tests/gpu/test_qwen3_5_deltanet.py``. """ @dataclass(kw_only=True, slots=True) class Config(Module.Config): - # "fla_chunked": parallel within chunks for training (default) - # "fla_fused_recurrent": for inference only in rl, no backward - backend: GatedDeltaBackend = "fla_chunked" + pass def __init__(self, config: Config): super().__init__() - self.backend = config.backend def forward( self, @@ -282,69 +228,36 @@ def forward( beta_TH: torch.Tensor, *, cu_seqlens: torch.Tensor | None = None, - cu_seqlens_cpu: torch.Tensor | None = None, ) -> torch.Tensor: - # Expand Q/K heads to match V when n_value_heads > n_key_heads - if xq_THK.shape[1] != xv_THV.shape[1]: - assert xv_THV.shape[1] % xq_THK.shape[1] == 0 - repeat = xv_THV.shape[1] // xq_THK.shape[1] - xq_THK = xq_THK.repeat_interleave(repeat, dim=1) - xk_THK = xk_THK.repeat_interleave(repeat, dim=1) - - xq_1THK = xq_THK.unsqueeze(0) - xk_1THK = xk_THK.unsqueeze(0) - xv_1THV = xv_THV.unsqueeze(0) - g_1TH = g_TH.unsqueeze(0) - beta_1TH = beta_TH.unsqueeze(0) + xq_BTHK = xq_THK.unsqueeze(0) + xk_BTHK = xk_THK.unsqueeze(0) + xv_BTHV = xv_THV.unsqueeze(0) + g_BTH = g_TH.unsqueeze(0) + beta_BTH = beta_TH.unsqueeze(0) if is_in_batch_invariant_mode() and cu_seqlens is not None: - if cu_seqlens_cpu is None: - raise ValueError( - "Batch-invariant Gated DeltaNet requires CPU cu_seqlens." - ) return _recurrent_gdn_fwd( - xq_1THK, - xk_1THK, - xv_1THV, - g_1TH, - beta_1TH, + xq_BTHK, + xk_BTHK, + xv_BTHV, + g_BTH, + beta_BTH, cu_seqlens, - cu_seqlens_cpu, ).squeeze(0) - if self.backend == "fla_chunked": - if cu_seqlens is not None and cu_seqlens_cpu is None: - raise ValueError( - "Qwen3.5 FLA varlen DeltaNet requires a CPU cu_seqlens tensor." - ) - result = _fla_chunk_gated_delta_rule( - xq_1THK, - xk_1THK, - xv_1THV, - g_1TH, - beta_1TH, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - elif self.backend == "fla_fused_recurrent": - result = _fla_fused_recurrent_gated_delta_rule( - xq_1THK, - xk_1THK, - xv_1THV, - g_1TH, - beta=beta_1TH, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) - else: - raise ValueError( - f"Unknown fla_backend '{self.backend}'. " - "Valid: 'fla_chunked', 'fla_fused_recurrent'." - ) - - # FLA kernels return (output, final_state); we only need output - return result[0].squeeze(0) + normalized_q = l2norm(xq_BTHK, cu_seqlens=cu_seqlens) + normalized_k = l2norm(xk_BTHK, cu_seqlens=cu_seqlens) + output, _ = chunk_gdn( + normalized_q, + normalized_k, + xv_BTHV, + g_BTH, + beta_BTH, + cu_seqlens=cu_seqlens, + scale=xq_BTHK.shape[-1] ** -0.5, + impl="fused", + ) + return output.squeeze(0) class InnerGatedDeltaNet(Module): @@ -379,30 +292,20 @@ def forward( *, key_head_dim: int, value_head_dim: int, - cu_seqlens_host: tuple[int, ...] | None = None, + use_packed_sequence: bool = False, ) -> torch.Tensor: """Run separate Q/K/V convolutions and recurrence on local heads.""" num_tokens = query_TC.shape[0] - if cu_seqlens_host is not None: - cu_seqlens_cpu = torch.tensor( - cu_seqlens_host, - dtype=cu_seqlens.dtype, - device="cpu", - ) - else: - cu_seqlens_cpu = None - def causal_conv( x_TC: torch.Tensor, weight_C1W: torch.Tensor, ) -> torch.Tensor: - if cu_seqlens_host is not None: + if use_packed_sequence: return _causal_conv1d_varlen( x_TC, weight_C1W, cu_seqlens, - cu_seqlens_cpu, ) x_1CT = F.pad( @@ -439,8 +342,7 @@ def causal_conv( xv_THV, g_TH, beta_TH, - cu_seqlens=cu_seqlens if cu_seqlens_host is not None else None, - cu_seqlens_cpu=cu_seqlens_cpu, + cu_seqlens=cu_seqlens if use_packed_sequence else None, ) @@ -507,19 +409,11 @@ def forward( attention_masks: VarlenMetadata | None = None, ) -> torch.Tensor: num_tokens = x_TD.shape[0] - cu_seqlens_host = None + use_packed_sequence = ( + attention_masks is not None or is_in_batch_invariant_mode() + ) if attention_masks is not None: - # FLA caches varlen index helpers by tensor identity. A fresh - # tensor ensures forward and activation-checkpoint recompute both - # execute the helpers instead of taking different cache paths. - with spmd.local(): - cu_seqlens = attention_masks.cu_seq_q.clone() - cu_seqlens_host = attention_masks.cu_seq_q_host - if cu_seqlens_host is None: - raise ValueError( - "Qwen3.5 GatedDeltaNet varlen requires CPU cu_seqlens " - "metadata. Build VarlenMetadata with include_host_offsets=True." - ) + cu_seqlens = attention_masks.cu_seq_q else: cu_seqlens = torch.arange( 0, @@ -528,8 +422,6 @@ def forward( dtype=torch.int32, device=x_TD.device, ) - if is_in_batch_invariant_mode(): - cu_seqlens_host = (0, num_tokens) query_TC = self.in_proj_q(x_TD) key_TC = self.in_proj_k(x_TD) @@ -552,7 +444,7 @@ def forward( cu_seqlens, key_head_dim=self.key_head_dim, value_head_dim=self.value_head_dim, - cu_seqlens_host=cu_seqlens_host, + use_packed_sequence=use_packed_sequence, ) gate_THV = gate_TC.view(num_tokens, -1, self.value_head_dim) output_THV = self.norm(output_THV, gate_THV) diff --git a/torchtitan/models/qwen3_5/sharding.py b/torchtitan/models/qwen3_5/sharding.py index ccbcda98a6..d4dd6e8dbc 100644 --- a/torchtitan/models/qwen3_5/sharding.py +++ b/torchtitan/models/qwen3_5/sharding.py @@ -12,7 +12,7 @@ Full-attention layers: TP on wq/wk/wv/wo with local_map for inner attention; each layer's MRoPE ``cache`` buffer is sharded Replicate. GatedDeltaNet layers: head-sharded TP on projections (ColwiseParallel) and -out_proj (RowwiseParallel); the FLA kernel and depthwise Conv1d run on local +out_proj (RowwiseParallel); the GDN kernel and depthwise Conv1d run on local tensors via local_map. """ @@ -360,7 +360,7 @@ def _set_deltanet_sharding( ) # The inner GDN is the DTensor-to-local boundary for the head-parallel - # convolution and recurrence. cu_seqlens_host is keyword-only host metadata + # convolution and recurrence. use_packed_sequence is keyword-only metadata # and intentionally remains outside local_map's positional placements. deltanet_cfg.inner_gated_delta_net.sharding_config = ShardingConfig( in_src_shardings={ diff --git a/torchtitan_recipes/tests/__init__.py b/torchtitan_recipes/tests/__init__.py index ecee18e423..0ffb42bebf 100644 --- a/torchtitan_recipes/tests/__init__.py +++ b/torchtitan_recipes/tests/__init__.py @@ -9,7 +9,7 @@ Each function here is one run of one entry in ``tests/integration_tests``, expressed as a full Trainer configuration. -Model registries that need an optional dependency (``fla``, ``torchvision``) or +Model registries that need an optional dependency (such as ``torchvision``) or that are slow to import are imported inside the function that uses them, so selecting any single configuration stays cheap. """ diff --git a/torchtitan_recipes/tests/models.py b/torchtitan_recipes/tests/models.py index d51ffe1522..e8e9135596 100644 --- a/torchtitan_recipes/tests/models.py +++ b/torchtitan_recipes/tests/models.py @@ -221,9 +221,6 @@ def qwen35_debugmodel_varlen_attn_fsdp2_tp2_sac() -> Trainer.Config: config = qwen35_debugmodel_varlen_attn(seq_len=512) config.parallelism.data_parallel_shard_degree = 2 config.parallelism.tensor_parallel_degree = 2 - # First-run FLA/TileLang kernel compile and autotune exceed the default - # 100s train timeout. - config.comm.train_timeout_seconds = 600 config.activation_checkpoint = SelectiveAC.Config() _use_spmd_types(config, typechecking=False) config.training.disable_cuda_graphs = True From d5ec7caba571797d78ad34a128b88a4d464bee78 Mon Sep 17 00:00:00 2001 From: drisspg Date: Tue, 1 Sep 2026 00:00:28 -0700 Subject: [PATCH 3/5] Use paged Attention Gym convolution for Qwen3.5 prefill ## Human Note ## Agent note Use Attention Gym's paged causal convolution API for vLLM multi-token prefill. The operation reads fresh or resumed cache slots directly and advances them in place, removing the remaining temporary history allocation, gather, host-syncing continuation check, and final-state scatter from Qwen3.5 serving. ## Performance On GB200 at the TP=2-local Qwen3.5 shape (C=4096, W=4, BF16), direct paging reduced convolution prefill GPU time by 43.8-79.4% for fresh prompts and 46.2-80.4% for resumed prefixes across N=1-64 and 128-512 tokens per sequence. ## Test Plan ```bash PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/attention-gym .venv/bin/pytest -q -n 6 test/test_short_conv_cute.py test/linear/test_gdn_chunk_fused.py test/test_namespaces.py PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=2 -m pytest -q torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen ``` --- torchtitan/experiments/rl/models/gdn.py | 50 ++++++------------------- 1 file changed, 11 insertions(+), 39 deletions(-) diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py index 4dc2c737d4..3941d205c2 100644 --- a/torchtitan/experiments/rl/models/gdn.py +++ b/torchtitan/experiments/rl/models/gdn.py @@ -16,9 +16,7 @@ cache stays in model dtype because it only stores trailing input columns. * Batch-invariant recurrence uses the same Attention Gym scan as the trainer. -Decode, recurrent execution, and chunked prefill all update the paged SSM state -pool directly. Convolution prefill still materializes its much smaller ``W - 1`` -history until Attention Gym exposes a paged varlen convolution operation. +Decode and prefill update the paged convolution and SSM state pools directly. """ from dataclasses import dataclass @@ -26,9 +24,9 @@ import torch import torch.nn.functional as F from attn_gym.linear import ( - causal_conv1d, causal_conv1d_decode, l2norm, + paged_causal_conv1d, paged_chunk_gdn, recurrent_gdn, recurrent_gdn_decode, @@ -277,11 +275,9 @@ def _forward( num_sequences = num_decodes + num_prefills # Convolution is split by request type and writes one contiguous - # conv_output for the single recurrence below. Decode is FULL-captured in - # a CUDA graph, so it must use the single-token update kernel: the varlen - # causal_conv1d prepares chunk indices with host syncs, which capture - # forbids. Prefill (eager at the graph break) uses the varlen kernel. - # vLLM orders tokens decode-first, then prefill. + # conv_output for the single recurrence below. Decode uses the specialized + # single-token state update, while prefill uses the packed multi-token + # operation. vLLM orders tokens decode-first, then prefill. conv_output = mixed_qkv.new_empty(num_actual_tokens, mixed_qkv.shape[1]) decode_slots = state_indices[:num_decodes] @@ -310,42 +306,18 @@ def _forward( # 0-based for the prefill slice. prefill_cu_seqlens = gdn_metadata.non_spec_query_start_loc else: - # Mixed batch: prefill_query_start_loc holds absolute offsets that - # start at num_decode_tokens, because decode tokens occupy the front - # of the batch. Subtract the first offset (which equals - # num_decode_tokens) to rebase the prefill slice's cu_seqlens to 0. + # Mixed-batch prefill metadata is already rebased to the prefill slice. assert gdn_metadata.prefill_query_start_loc is not None - prefill_cu_seqlens = ( - gdn_metadata.prefill_query_start_loc - - gdn_metadata.prefill_query_start_loc[0] - ) - num_prefill_sequences = int(prefill_cu_seqlens.numel()) - 1 - # This implementation runs eager at the graph break, so checking - # whether any prefix state must be restored does not enter a captured graph. - has_continuations = prefill_has_initial_state is not None and bool( - prefill_has_initial_state.any() - ) - conv_initial_state = mixed_qkv.new_zeros( - num_prefill_sequences, - self.conv_kernel_size - 1, - mixed_qkv.shape[1], - ) - # Fresh prefills keep zero state; prefix-cache continuations restore - # only the sequence slots identified by vLLM metadata. - if has_continuations: - resumed_slots = prefill_slots[prefill_has_initial_state] - conv_initial_state[prefill_has_initial_state] = conv_state[ - resumed_slots - ] - prefill_conv_output, conv_final_state = causal_conv1d( + prefill_cu_seqlens = gdn_metadata.prefill_query_start_loc + prefill_conv_output = paged_causal_conv1d( mixed_qkv[prefill_start:num_actual_tokens].unsqueeze(0), conv_weight, + conv_state, + prefill_slots, activation="silu", cu_seqlens=prefill_cu_seqlens, - initial_state=conv_initial_state, - return_final_state=True, + has_initial_state=prefill_has_initial_state, ) - conv_state[prefill_slots] = conv_final_state.to(conv_state.dtype) conv_output[prefill_start:num_actual_tokens] = prefill_conv_output.squeeze( 0 ) From 1d6d9cb7316b744dcf3c5ddf719db26a4e1bd672 Mon Sep 17 00:00:00 2001 From: drisspg Date: Tue, 1 Sep 2026 09:55:18 -0700 Subject: [PATCH 4/5] Address Qwen3.5 Attention Gym review feedback ## Human Note ## Agent note Remove the redundant subprocess import test, narrow the recurrent backward smoke check to the decay gradient it is intended to protect, and use the current expected vLLM batch-invariance import path directly instead of carrying compatibility logic for an older package layout. ## Test Plan ```bash PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/torchtitan python -c 'from torchtitan.experiments.rl.batch_invariance import patch_bmm_for_batch_invariance; patch_bmm_for_batch_invariance()' ``` --- tests/unit_tests/gpu/test_qwen3_5_deltanet.py | 5 ++--- tests/unit_tests/test_qwen3_5_mrope_positions.py | 15 --------------- torchtitan/experiments/rl/batch_invariance.py | 10 +--------- 3 files changed, 3 insertions(+), 27 deletions(-) diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index b020a32a7d..0af78908de 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -697,9 +697,8 @@ def test_batch_invariant_recurrent_matches_paged_attention_gym(self): torch.testing.assert_close(actual, expected, rtol=0, atol=0) actual.float().square().mean().backward() - for tensor in (q, k, v, decay, update_gate): - self.assertIsNotNone(tensor.grad) - self.assertTrue(torch.isfinite(tensor.grad).all()) + self.assertIsNotNone(decay.grad) + self.assertTrue(torch.isfinite(decay.grad).all()) if __name__ == "__main__": diff --git a/tests/unit_tests/test_qwen3_5_mrope_positions.py b/tests/unit_tests/test_qwen3_5_mrope_positions.py index 4ab1c8a988..2a994f1546 100644 --- a/tests/unit_tests/test_qwen3_5_mrope_positions.py +++ b/tests/unit_tests/test_qwen3_5_mrope_positions.py @@ -18,8 +18,6 @@ resolution lives in ``forward`` or in ``preprocess_inputs``. """ -import subprocess -import sys import unittest import torch @@ -56,19 +54,6 @@ def forward(self, x, attention_masks=None, positions=None): class TestQwen35MRoPEPositions(unittest.TestCase): - def test_model_import_does_not_require_fla(self): - script = ( - "import builtins\n" - "original_import = builtins.__import__\n" - "def without_fla(name, globals=None, locals=None, fromlist=(), level=0):\n" - " if level == 0 and (name == 'fla' or name.startswith('fla.')):\n" - " raise ModuleNotFoundError('blocked fla import')\n" - " return original_import(name, globals, locals, fromlist, level)\n" - "builtins.__import__ = without_fla\n" - "import torchtitan.models.qwen3_5\n" - ) - subprocess.run([sys.executable, "-c", script], check=True) - def _build_stub_model(self): model_registry, ParallelDims, ParallelismConfig = _build_config_modules() # varlen backend keeps mask construction to pure tensor ops (no flex diff --git a/torchtitan/experiments/rl/batch_invariance.py b/torchtitan/experiments/rl/batch_invariance.py index 20ce18cda5..3dd6f5bfb8 100644 --- a/torchtitan/experiments/rl/batch_invariance.py +++ b/torchtitan/experiments/rl/batch_invariance.py @@ -72,15 +72,7 @@ def patch_bmm_for_batch_invariance() -> None: global _batch_invariant_bmm_lib if _batch_invariant_bmm_lib is not None: return - try: - from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant - except ModuleNotFoundError as error: - if error.name not in { - "vllm.model_executor.determinism", - "vllm.model_executor.determinism.batch_invariant", - }: - raise - from vllm.model_executor.layers.batch_invariant import bmm_batch_invariant + from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant _batch_invariant_bmm_lib = torch.library.Library("aten", "IMPL") _batch_invariant_bmm_lib.impl("bmm", bmm_batch_invariant, "CUDA") From 445ccc32079ee9635315e7c0b9726bdcb48ad991 Mon Sep 17 00:00:00 2001 From: drisspg Date: Tue, 1 Sep 2026 17:53:10 -0700 Subject: [PATCH 5/5] Clarify Qwen3.5 GDN dispatch semantics --- .ci/docker/requirements.txt | 2 +- pyproject.toml | 2 +- tests/unit_tests/cpu/test_qwen3_8.py | 2 +- tests/unit_tests/cpu/test_varlen_attention.py | 6 +- tests/unit_tests/gpu/test_qwen3_5_deltanet.py | 38 +++--- tests/unit_tests/test_kda_attention.py | 6 +- .../test_qwen3_5_mrope_positions.py | 10 +- torchtitan/experiments/rl/models/gdn.py | 117 ++++++++---------- torchtitan/models/common/attention.py | 25 +--- torchtitan/models/kimi_k3/README.md | 2 +- torchtitan/models/kimi_k3/kda.py | 1 + torchtitan/models/qwen3_5/README.md | 2 +- torchtitan/models/qwen3_5/__init__.py | 1 + torchtitan/models/qwen3_5/gdn.py | 10 +- torchtitan/models/qwen3_5/model.py | 15 +-- torchtitan/models/qwen3_5/sharding.py | 3 +- 16 files changed, 93 insertions(+), 149 deletions(-) diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt index 7fb150c9ac..e47519750b 100644 --- a/.ci/docker/requirements.txt +++ b/.ci/docker/requirements.txt @@ -8,4 +8,4 @@ safetensors einops pillow spmd_types==0.2.5 -attn-gym[linear]==0.0.6 +attn-gym[linear]==0.0.8 diff --git a/pyproject.toml b/pyproject.toml index 41b673d48e..b0cf862749 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "einops", "pillow", "spmd_types==0.2.5", - "attn-gym[linear]==0.0.6", + "attn-gym[linear]==0.0.8", ] dynamic = ["version"] diff --git a/tests/unit_tests/cpu/test_qwen3_8.py b/tests/unit_tests/cpu/test_qwen3_8.py index 4b766643e4..8cfbdf83b9 100644 --- a/tests/unit_tests/cpu/test_qwen3_8.py +++ b/tests/unit_tests/cpu/test_qwen3_8.py @@ -10,7 +10,7 @@ import pytest import torch -pytest.importorskip("fla") +pytest.importorskip("attn_gym") from torchtitan.models.qwen3_5 import Qwen35Model, Qwen35StateDictAdapter from torchtitan.models.qwen3_5.sharding import set_qwen35_sharding_config diff --git a/tests/unit_tests/cpu/test_varlen_attention.py b/tests/unit_tests/cpu/test_varlen_attention.py index e23b0b0d2d..8aa30c5cd5 100644 --- a/tests/unit_tests/cpu/test_varlen_attention.py +++ b/tests/unit_tests/cpu/test_varlen_attention.py @@ -29,17 +29,13 @@ class TestPackedVarlenMetadata(unittest.TestCase): def test_document_boundaries(self): positions_T = torch.tensor([0, 1, 2, 0, 1, 0, 1, 2, 3]) - metadata = create_varlen_metadata_for_document( - positions_T, - include_host_offsets=True, - ) + metadata = create_varlen_metadata_for_document(positions_T) expected_cu_seq = torch.tensor([0, 3, 5, 9], dtype=torch.int32) torch.testing.assert_close(metadata.cu_seq_q, expected_cu_seq) torch.testing.assert_close(metadata.cu_seq_k, expected_cu_seq) self.assertEqual(metadata.max_q, 4) self.assertEqual(metadata.max_k, 4) - self.assertEqual(metadata.cu_seq_q_host, (0, 3, 5, 9)) class TestPackedVarlenAttention(unittest.TestCase): diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index 0af78908de..a6786ef80c 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -177,7 +177,8 @@ def test_flex_masks_ignore_padding_position_resets(self): ) from exc with torch.device("meta"): - model = qwen3_5_configs["debugmodel"]("flex").build() + build_config, max_context_length = qwen3_5_configs["debugmodel"] + model = build_config("flex", seq_len=max_context_length).build() positions = torch.tensor([0, 1, 2, 0, 0], dtype=torch.int32) with mock.patch.object(Decoder, "get_attention_masks", return_value=None): @@ -195,7 +196,8 @@ def test_flex_masks_include_delta_net_varlen_metadata(self): ) from exc with torch.device("meta"): - model = qwen3_5_configs["debugmodel"]("flex").build() + build_config, max_context_length = qwen3_5_configs["debugmodel"] + model = build_config("flex", seq_len=max_context_length).build() positions = torch.tensor([0, 1, 0, 1, 2], dtype=torch.int32) full_attention_mask = mock.sentinel.full_attention_mask @@ -211,7 +213,6 @@ def test_flex_masks_include_delta_net_varlen_metadata(self): attention_masks["deltanet"].cu_seq_q, torch.tensor([0, 2, 5], dtype=torch.int32), ) - self.assertEqual(attention_masks["deltanet"].cu_seq_q_host, (0, 2, 5)) def _make_deltanet( self, @@ -366,10 +367,7 @@ def test_extracted_forward_matches_main(self): [0, 1, 0, 1, 2, 0, 1, 2, 0, 1], dtype=torch.int32, ) - attention_masks = create_varlen_metadata_for_document( - positions, - include_host_offsets=True, - ) + attention_masks = create_varlen_metadata_for_document(positions) for masks in (None, attention_masks): with mock.patch( @@ -410,10 +408,7 @@ def test_varlen_matches_independent_document_forwards(self): dtype=torch.int32, ) - attention_masks = create_varlen_metadata_for_document( - positions, - include_host_offsets=False, - ) + attention_masks = create_varlen_metadata_for_document(positions) self._assert_packed_run_matches_per_document( model, x_TD, positions, attention_masks ) @@ -447,7 +442,10 @@ def test_get_attention_masks_pairs_flex_mask_with_deltanet_offsets(self): self.assertIsInstance(masks["quadratic_attention"], BlockMask) self.assertIsInstance(masks["deltanet"], VarlenMetadata) # Three packed documents have lengths 3, 2, and 5. - self.assertEqual(masks["deltanet"].cu_seq_q_host, (0, 3, 5, 10)) + torch.testing.assert_close( + masks["deltanet"].cu_seq_q, + torch.tensor([0, 3, 5, 10], dtype=torch.int32, device=device), + ) # Each block picks the entry matching its layer type. mask_keys = {layer.attn_mask_key for layer in flex_model.layers.values()} @@ -463,7 +461,10 @@ def test_get_attention_masks_pairs_flex_mask_with_deltanet_offsets(self): self.assertIsInstance(varlen_masks, dict) self.assertIs(varlen_masks["quadratic_attention"], varlen_masks["deltanet"]) self.assertIsInstance(varlen_masks["deltanet"], VarlenMetadata) - self.assertEqual(varlen_masks["deltanet"].cu_seq_q_host, (0, 3, 5, 10)) + torch.testing.assert_close( + varlen_masks["deltanet"].cu_seq_q, + torch.tensor([0, 3, 5, 10], dtype=torch.int32, device=device), + ) deltanet_only_config = model_registry("debugmodel").model deltanet_only_config.layers = [ @@ -479,9 +480,9 @@ def test_get_attention_masks_pairs_flex_mask_with_deltanet_offsets(self): ) self.assertIsNone(deltanet_only_masks["quadratic_attention"]) self.assertIsInstance(deltanet_only_masks["deltanet"], VarlenMetadata) - self.assertEqual( - deltanet_only_masks["deltanet"].cu_seq_q_host, - (0, 3, 5, 10), + torch.testing.assert_close( + deltanet_only_masks["deltanet"].cu_seq_q, + torch.tensor([0, 3, 5, 10], dtype=torch.int32, device=device), ) def _assert_fused_varlen_matches_per_document( @@ -547,10 +548,7 @@ def _assert_fused_varlen_matches_per_document( requires_grad=True, ) - attention_masks = create_varlen_metadata_for_document( - positions, - include_host_offsets=True, - ) + attention_masks = create_varlen_metadata_for_document(positions) actual = model(x_TD, attention_masks) # Reference: run each document on its own and stitch the outputs back. diff --git a/tests/unit_tests/test_kda_attention.py b/tests/unit_tests/test_kda_attention.py index 015871915b..ac0060664f 100644 --- a/tests/unit_tests/test_kda_attention.py +++ b/tests/unit_tests/test_kda_attention.py @@ -92,11 +92,7 @@ def test_varlen_matches_independent_documents(self): device="cuda", dtype=torch.int32, ) - masks = create_varlen_metadata_for_document( - positions_T, - include_host_offsets=True, - ) - self.assertEqual(masks.cu_seq_q_host, (0, 37, 101, 192)) + masks = create_varlen_metadata_for_document(positions_T) model = self._make_kda() packed_TD = model(x_TD, masks) diff --git a/tests/unit_tests/test_qwen3_5_mrope_positions.py b/tests/unit_tests/test_qwen3_5_mrope_positions.py index 2a994f1546..2a3208d605 100644 --- a/tests/unit_tests/test_qwen3_5_mrope_positions.py +++ b/tests/unit_tests/test_qwen3_5_mrope_positions.py @@ -94,8 +94,9 @@ def test_text_batch_routes_1d_positions_to_layers(self): # No mrope: layers see the plain 1D positions. self.assertTrue(torch.equal(sink["positions"], positions)) # Masks come from the 1D positions. - self.assertEqual( - batch["attention_masks"]["deltanet"].cu_seq_q_host, (0, 3, 5, 10) + torch.testing.assert_close( + batch["attention_masks"]["deltanet"].cu_seq_q, + torch.tensor([0, 3, 5, 10], dtype=torch.int32, device=positions.device), ) def test_multimodal_batch_routes_mrope_to_layers(self): @@ -122,8 +123,9 @@ def test_multimodal_batch_routes_mrope_to_layers(self): self.assertEqual(sink["positions"].shape[-1], 3) self.assertTrue(torch.equal(sink["positions"], mrope_positions)) # Masks are still built from the 1D positions, not the mrope positions. - self.assertEqual( - batch["attention_masks"]["deltanet"].cu_seq_q_host, (0, 3, 5, 10) + torch.testing.assert_close( + batch["attention_masks"]["deltanet"].cu_seq_q, + torch.tensor([0, 3, 5, 10], dtype=torch.int32, device=positions.device), ) diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py index 3941d205c2..07d13d3260 100644 --- a/torchtitan/experiments/rl/models/gdn.py +++ b/torchtitan/experiments/rl/models/gdn.py @@ -80,6 +80,8 @@ def __init__(self, config: Config) -> None: self.num_speculative_tokens = ( speculative_config.num_speculative_tokens if speculative_config else 0 ) + # vLLM speculative decoding retains one state per draft position and + # commits the accepted state; Attention Gym currently mutates one slot. if self.num_speculative_tokens != 0: raise ValueError("Attention Gym GDN does not support speculative decoding.") @@ -110,6 +112,13 @@ def __init__(self, config: Config) -> None: "the paged convolution history has contiguous channels." ) + # Attention Gym's paged kernels mutate the SSM state pool directly and + # require FP32 state in both regular and batch-invariant execution. + if self.cache_config.mamba_ssm_cache_dtype not in {"auto", "float32"}: + raise ValueError( + "Attention Gym GDN requires mamba_ssm_cache_dtype='float32', " + f"got {self.cache_config.mamba_ssm_cache_dtype!r}." + ) self.cache_config.mamba_ssm_cache_dtype = "float32" # vLLM populates this via the KV-cache allocator: (conv_state, ssm_state). @@ -176,56 +185,6 @@ def _split_qkv( ) return query, key, value - def _run_recurrence( - self, - conv_output: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - negative_exp_A: torch.Tensor, - dt_bias: torch.Tensor, - ssm_state: torch.Tensor, - slot_indices: torch.Tensor, - cu_seqlens: torch.Tensor, - has_initial_state: torch.Tensor | None, - batch_invariant: bool, - ) -> torch.Tensor: - """Run the selected recurrence and update the paged SSM slots.""" - query, key, value = self._split_qkv(conv_output) - decay = (negative_exp_A * F.softplus(a.float() + dt_bias)).unsqueeze(0) - update_gate = torch.sigmoid(b).unsqueeze(0) - query = l2norm(query, cu_seqlens=cu_seqlens) - key = l2norm(key, cu_seqlens=cu_seqlens) - - if batch_invariant: - output, _ = recurrent_gdn( - query, - key, - value, - decay, - update_gate, - ssm_state, - cu_seqlens=cu_seqlens, - scale=self.head_k_dim**-0.5, - state_indices=slot_indices, - has_initial_state=has_initial_state, - # Triton autotuning breaks batch invariance. - autotune=False, - ) - return output - - return paged_chunk_gdn( - query, - key, - value, - decay, - update_gate, - ssm_state, - slot_indices, - cu_seqlens=cu_seqlens, - has_initial_state=has_initial_state, - scale=self.head_k_dim**-0.5, - ) - # The decorator makes this an eager graph-split point during breakable capture. # The caller-owned output has a stable address across graph replays. @eager_break_during_capture @@ -267,7 +226,6 @@ def _forward( ssm_state = self.kv_cache[1] conv_state = self.kv_cache[0] assert conv_bias is None - negative_exp_A = -torch.exp(A_log.float()) dt_bias = dt_bias.float() num_decodes = gdn_metadata.num_decodes num_prefills = gdn_metadata.num_prefills @@ -297,7 +255,9 @@ def _forward( assert gdn_metadata.prefill_state_indices is not None prefill_slots = gdn_metadata.prefill_state_indices prefill_has_initial_state = gdn_metadata.prefill_has_initial_state - assert prefill_has_initial_state is not None + assert ( + prefill_has_initial_state is not None + ), "prefill_has_initial_state is required when num_prefills > 0" prefill_start = num_decode_tokens if num_decodes > 0 else 0 # cu_seqlens must be 0-based within the prefill slice that the conv # kernel receives (mixed_qkv[prefill_start:]). @@ -337,11 +297,12 @@ def _forward( has_initial_state = torch.ones( num_sequences, dtype=torch.bool, device=mixed_qkv.device ) - if prefill_has_initial_state is not None: - has_initial_state[num_decodes:] = prefill_has_initial_state + has_initial_state[num_decodes:] = prefill_has_initial_state batch_invariant = is_in_batch_invariant_mode() if num_prefills == 0 and not batch_invariant: + # Pure decode fuses QKV splitting, gate activation, normalization, + # and the recurrent update into one kernel. recurrent_gdn_decode( conv_output[:num_decode_tokens], a[:num_decode_tokens].unsqueeze(0), @@ -355,18 +316,42 @@ def _forward( ) return - recurrent_output = self._run_recurrence( - conv_output, - a, - b, - negative_exp_A, - dt_bias, - ssm_state, - all_slots, - cu_seqlens, - has_initial_state, - batch_invariant=batch_invariant, + query, key, value = self._split_qkv(conv_output) + decay = (-torch.exp(A_log.float()) * F.softplus(a.float() + dt_bias)).unsqueeze( + 0 ) + update_gate = torch.sigmoid(b).unsqueeze(0) + query = l2norm(query, cu_seqlens=cu_seqlens) + key = l2norm(key, cu_seqlens=cu_seqlens) + + if batch_invariant: + recurrent_output, _ = recurrent_gdn( + query, + key, + value, + decay, + update_gate, + ssm_state, + cu_seqlens=cu_seqlens, + scale=self.head_k_dim**-0.5, + state_indices=all_slots, + has_initial_state=has_initial_state, + # Triton autotuning breaks batch invariance. + autotune=False, + ) + else: + recurrent_output = paged_chunk_gdn( + query, + key, + value, + decay, + update_gate, + ssm_state, + all_slots, + cu_seqlens=cu_seqlens, + has_initial_state=has_initial_state, + scale=self.head_k_dim**-0.5, + ) output[:num_actual_tokens] = recurrent_output[0, :num_actual_tokens].to( output.dtype ) @@ -387,7 +372,7 @@ def forward( *, key_head_dim: int, value_head_dim: int, - use_packed_sequence: bool = False, + use_varlen_kernels: bool = False, ) -> torch.Tensor: """Run the flattened vLLM cache operation on rank-local tensors.""" assert key_head_dim == self.head_k_dim diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index 941e3f0589..28a74db94f 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -80,7 +80,6 @@ class VarlenMetadata(NamedTuple): cu_seq_k: torch.Tensor max_q: int max_k: int - cu_seq_q_host: tuple[int, ...] | None = None # Mapping (not dict) lets covariant value types accept both BlockMask-only @@ -592,8 +591,6 @@ def create_attention_mask(*args, **kwargs): def create_varlen_metadata_for_document( positions: torch.Tensor, - *, - include_host_offsets: bool = False, ) -> VarlenMetadata: """Creates cumulative sequence length indices needed for variable length attention. @@ -603,8 +600,6 @@ def create_varlen_metadata_for_document( Args: positions: Per-token position tensor with shape ``[T]``. Positions reset to 0 at each document start. - include_host_offsets: Also materialize cumulative sequence offsets as - host metadata for kernels that need it. Returns: VarlenMetadata containing cumulative sequence length indices for q, k, @@ -625,24 +620,7 @@ def create_varlen_metadata_for_document( spmd.mutate_type(packed_cu_seqlens, "dp", src=spmd.R, dst=spmd.V) seq_lengths = torch.diff(packed_cu_seqlens) - max_seqlen: int - packed_cu_seqlens_host = None - if include_host_offsets: - packed_cu_seqlens_host = tuple( - int(offset) for offset in packed_cu_seqlens.tolist() - ) - max_seqlen = max( - ( - end - start - for start, end in zip( - packed_cu_seqlens_host[:-1], - packed_cu_seqlens_host[1:], - strict=False, - ) - ), - default=0, - ) - elif seq_lengths.numel() > 0: + if seq_lengths.numel() > 0: # device to host sync but only done once per model forward max_seqlen = int(seq_lengths.max().item()) else: @@ -653,7 +631,6 @@ def create_varlen_metadata_for_document( cu_seq_k=packed_cu_seqlens, max_q=max_seqlen, max_k=max_seqlen, - cu_seq_q_host=packed_cu_seqlens_host, ) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 765d72825d..b2304552a1 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -8,7 +8,7 @@ Attention (MLA) decoder with LatentMoE and a MoonViT-V2 vision encoder. Install the additional dependencies: ```bash -pip install av einops pillow torchvision +pip install -r .ci/docker/requirements-vlm.txt ``` ## Architecture diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index 552d7929e4..fe5dfcf947 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -96,6 +96,7 @@ def forward( A_log_H.float(), dt_bias_HK.float(), lower_bound=self.lower_bound, + impl="fused", ) output_1THV, _ = chunk_kda( l2norm(q_1THK), diff --git a/torchtitan/models/qwen3_5/README.md b/torchtitan/models/qwen3_5/README.md index 2147416772..c336a4e1cc 100644 --- a/torchtitan/models/qwen3_5/README.md +++ b/torchtitan/models/qwen3_5/README.md @@ -24,7 +24,7 @@ Note: the diagram shows each patch mapping to one vision token. In practice, the Install the additional dependencies: ```bash -pip install av torchvision +pip install -r .ci/docker/requirements-vlm.txt ``` ## Model Variants diff --git a/torchtitan/models/qwen3_5/__init__.py b/torchtitan/models/qwen3_5/__init__.py index 34bffdd89e..642524d581 100644 --- a/torchtitan/models/qwen3_5/__init__.py +++ b/torchtitan/models/qwen3_5/__init__.py @@ -507,6 +507,7 @@ def _debugmodel(attn_backend: str, *, seq_len: int) -> Qwen35Model.Config: hidden_dim=512, n_key_heads=2, n_value_heads=4, + # Attention Gym fused chunk GDN requires K=V=128 on SM80+. key_head_dim=128, value_head_dim=128, ), diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py index daba81b9ba..6223576e6d 100644 --- a/torchtitan/models/qwen3_5/gdn.py +++ b/torchtitan/models/qwen3_5/gdn.py @@ -292,16 +292,16 @@ def forward( *, key_head_dim: int, value_head_dim: int, - use_packed_sequence: bool = False, ) -> torch.Tensor: """Run separate Q/K/V convolutions and recurrence on local heads.""" num_tokens = query_TC.shape[0] + use_varlen_kernels = cu_seqlens.numel() > 2 or is_in_batch_invariant_mode() def causal_conv( x_TC: torch.Tensor, weight_C1W: torch.Tensor, ) -> torch.Tensor: - if use_packed_sequence: + if use_varlen_kernels: return _causal_conv1d_varlen( x_TC, weight_C1W, @@ -342,7 +342,7 @@ def causal_conv( xv_THV, g_TH, beta_TH, - cu_seqlens=cu_seqlens if use_packed_sequence else None, + cu_seqlens=cu_seqlens if use_varlen_kernels else None, ) @@ -409,9 +409,6 @@ def forward( attention_masks: VarlenMetadata | None = None, ) -> torch.Tensor: num_tokens = x_TD.shape[0] - use_packed_sequence = ( - attention_masks is not None or is_in_batch_invariant_mode() - ) if attention_masks is not None: cu_seqlens = attention_masks.cu_seq_q else: @@ -444,7 +441,6 @@ def forward( cu_seqlens, key_head_dim=self.key_head_dim, value_head_dim=self.value_head_dim, - use_packed_sequence=use_packed_sequence, ) gate_THV = gate_TC.view(num_tokens, -1, self.value_head_dim) output_THV = self.norm(output_THV, gate_THV) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index d273f7e496..1a88c714cc 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -461,17 +461,10 @@ def get_attention_masks( first_token = torch.arange(positions.shape[0], device=positions.device) == 0 sequence_starts = ((positions == 0) & followed_by_one) | first_token sequence_positions = torch.where(sequence_starts, 0, 1) - deltanet_metadata = create_varlen_metadata_for_document( - sequence_positions, - include_host_offsets=True, - ) - if ( - deltanet_metadata.cu_seq_q_host is not None - and len(deltanet_metadata.cu_seq_q_host) == 2 - and not ( - attn_config is not None - and isinstance(attn_config.inner_attention, VarlenAttention.Config) - ) + deltanet_metadata = create_varlen_metadata_for_document(sequence_positions) + if deltanet_metadata.cu_seq_q.numel() == 2 and not ( + attn_config is not None + and isinstance(attn_config.inner_attention, VarlenAttention.Config) ): deltanet_metadata = None diff --git a/torchtitan/models/qwen3_5/sharding.py b/torchtitan/models/qwen3_5/sharding.py index d4dd6e8dbc..119fa5f5e2 100644 --- a/torchtitan/models/qwen3_5/sharding.py +++ b/torchtitan/models/qwen3_5/sharding.py @@ -360,8 +360,7 @@ def _set_deltanet_sharding( ) # The inner GDN is the DTensor-to-local boundary for the head-parallel - # convolution and recurrence. use_packed_sequence is keyword-only metadata - # and intentionally remains outside local_map's positional placements. + # convolution and recurrence. deltanet_cfg.inner_gated_delta_net.sharding_config = ShardingConfig( in_src_shardings={ "query_TC": projected_placement,