Cleanup readme and remove unneeded flashv3 backward

Author: bwasti

torchtitan/experiments/deterministic_vllm_rl/README.md

@@ -226,6 +226,10 @@ deterministic_vllm_rl/
 │   └── qwen3/
 │       ├── __init__.py
 │       └── model_vllm_compat.py       # vLLM-compatible Qwen3 model
+├── weights/
+│   ├── __init__.py
+│   ├── converter.py                   # Weight conversion script
+│   └── README.md                      # Weight conversion documentation
 └── tests/
     ├── __init__.py
     ├── test_batch_invariant_backward.py  # Test backward passes

torchtitan/experiments/deterministic_vllm_rl/batch_invariant_backward.py

@@ -42,109 +42,6 @@
 # Custom autograd Functions for vLLM operations
 # ============================================================================
 
-class FlashAttn3Function(Function):
-    """
-    Autograd function for Flash Attention 3 with proper backward support.
-    """
-
-    @staticmethod
-    def forward(
-        ctx,
-        q, k, v,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        softmax_scale,
-        causal,
-        window_left,
-        window_right,
-        softcap,
-        scheduler_metadata,
-        num_splits,
-    ):
-        """
-        Forward pass using vLLM's FA3 CUDA kernel.
-        """
-        out, softmax_lse, _, _ = torch.ops._vllm_fa3_C.fwd(
-            q, k, v,
-            None, None,       # k_new, v_new
-            None,             # q_v
-            None,             # out
-            cu_seqlens_q,
-            cu_seqlens_k,
-            None,             # cu_seqlens_k_new
-            None, None,       # seqused_q, seqused_k
-            max_seqlen_q, max_seqlen_k,
-            None,             # block_table
-            None,             # kv_batch_idx
-            None,             # leftpad_k
-            None, None, None, # rotary_cos, rotary_sin, seqlens_rotary
-            None, None, None, # q_descale, k_descale, v_descale
-            softmax_scale,
-            causal,
-            window_left, window_right,
-            softcap,
-            True,             # rotary_interleaved
-            scheduler_metadata,
-        )
-
-        # Save tensors needed for backward
-        ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k)
-        ctx.softmax_scale = softmax_scale
-        ctx.causal = causal
-        ctx.window_left = window_left
-        ctx.window_right = window_right
-        ctx.softcap = softcap
-        ctx.max_seqlen_q = max_seqlen_q
-        ctx.max_seqlen_k = max_seqlen_k
-        ctx.scheduler_metadata = scheduler_metadata
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        """
-        Backward pass using vLLM's FA3 CUDA backward kernel.
-        """
-        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
-
-        # Allocate gradient tensors
-        grad_q = torch.empty_like(q)
-        grad_k = torch.empty_like(k)
-        grad_v = torch.empty_like(v)
-
-        # Call FA3 backward kernel
-        torch.ops._vllm_fa3_C.bwd(
-            grad_output,
-            q, k, v,
-            out,
-            softmax_lse,
-            grad_q, grad_k, grad_v,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            None,             # cu_seqlens_k_new
-            None, None,       # seqused_q, seqused_k
-            ctx.max_seqlen_q,
-            ctx.max_seqlen_k,
-            None,             # block_table
-            None,             # kv_batch_idx
-            None,             # leftpad_k
-            None, None, None, # rotary_cos, rotary_sin, seqlens_rotary
-            None, None, None, # dq_accum, q_descale, k_descale, v_descale
-            ctx.softmax_scale,
-            ctx.causal,
-            ctx.window_left, ctx.window_right,
-            ctx.softcap,
-            False,            # deterministic
-            True,             # rotary_interleaved
-            ctx.scheduler_metadata,
-        )
-
-        # Return gradients for all forward inputs (None for non-tensor args)
-        return grad_q, grad_k, grad_v, None, None, None, None, None, None, None, None, None, None, None
-
-
 class SiluAndMulFunction(Function):
     """
     Autograd function for vLLM's SiluAndMul activation.
@@ -459,48 +356,6 @@ def patch_batch_invariant_with_gradients():
     _batch_invariant_backward_LIB.impl("aten::matmul_backward", matmul_backward_impl, "CUDA")
     _batch_invariant_backward_LIB.impl("aten::linear_backward", linear_backward_impl, "CUDA")
 
-    # Monkey-patch vLLM's flash_attn_varlen_func to use our autograd wrapper for FA3
-    import vllm.vllm_flash_attn.flash_attn_interface as fa_interface
-    _original_flash_attn_varlen_func = fa_interface.flash_attn_varlen_func
-
-    def patched_flash_attn_varlen_func(*args, **kwargs):
-        # Only patch FA3 calls
-        fa_version = kwargs.get('fa_version', fa_interface.DEFAULT_FA_VERSION)
-        if fa_version == 3:
-            # Extract the args needed for our autograd function
-            q = args[0]
-            k = args[1]
-            v = args[2]
-            max_seqlen_q = args[3]
-            cu_seqlens_q = args[4]
-            max_seqlen_k = args[5]
-            cu_seqlens_k = args[6] if len(args) > 6 else kwargs.get('cu_seqlens_k')
-            softmax_scale = kwargs.get('softmax_scale')
-            causal = kwargs.get('causal', False)
-            window_size = kwargs.get('window_size', (-1, -1))
-            softcap = kwargs.get('softcap', 0.0)
-            scheduler_metadata = kwargs.get('scheduler_metadata')
-            num_splits = kwargs.get('num_splits', 0)
-
-            if window_size is None:
-                window_size = (-1, -1)
-            window_left, window_right = window_size
-
-            # Use our autograd wrapper
-            return FlashAttn3Function.apply(
-                q, k, v,
-                cu_seqlens_q, cu_seqlens_k,
-                max_seqlen_q, max_seqlen_k,
-                softmax_scale, causal,
-                window_left, window_right,
-                softcap, scheduler_metadata, num_splits
-            )
-        else:
-            # Fall through to original implementation for FA2
-            return _original_flash_attn_varlen_func(*args, **kwargs)
-
-    fa_interface.flash_attn_varlen_func = patched_flash_attn_varlen_func
-
     _batch_invariant_backward_MODE = True