lint

Author: Gasoonjia

backends/cuda/runtime/cuda_backend.cpp

@@ -262,6 +262,27 @@ class ET_EXPERIMENTAL CudaBackend final
       FreeableBuffer* processed, // This will be a empty buffer
       ArrayRef<CompileSpec> compile_specs // This will be my empty list
   ) const override {
+    // Apply runtime_specs from LoadBackendOptionsMap (if provided)
+    auto runtime_specs = context.runtime_specs();
+    if (runtime_specs.size() > 0) {
+      for (size_t i = 0; i < runtime_specs.size(); ++i) {
+        const auto& opt = runtime_specs[i];
+        if (std::strcmp(opt.key, kSkipCopyOutputToCpuForMethod) == 0) {
+          if (auto* val =
+                  std::get_if<std::array<char, kMaxOptionValueLength>>(
+                      &opt.value)) {
+            const_cast<CudaBackend*>(this)->set_skip_copy_method(*val);
+          }
+        } else if (std::strcmp(opt.key, kUseSharedCudaStream) == 0) {
+          if (auto* val = std::get_if<bool>(&opt.value)) {
+            if (*val) {
+              const_cast<CudaBackend*>(this)->create_shared_cuda_stream();
+            }
+          }
+        }
+      }
+    }
+
     std::string method_name;
     for (const CompileSpec& spec : compile_specs) {
       if (std::strcmp(spec.key, "method_name") == 0) {

backends/cuda/tests/test_chunk_gated_delta_rule.py

@@ -42,7 +42,8 @@
 from torch.export import export
 
 
-B, T, H, K, V = 1, 128, 4, 64, 64
+B, H, K, V = 1, 4, 64, 64
+T = 128  # default T for chunked tests
 
 EXECUTORCH_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "../../.."))
 RUNNER_PATH = os.path.join(EXECUTORCH_ROOT, "cmake-out", "executor_runner")
@@ -88,32 +89,37 @@ def _make_inputs_from_fla(
     gate_logit_normalizer,
     mask_p=0.0,
     nonzero_h0=False,
+    seq_len=T,
     dtype=torch.bfloat16,
     device="cuda",
 ):
     """Generate inputs following FLA test_chunk() conventions."""
     torch.manual_seed(seed)
-    q = torch.rand(B, T, H, K, dtype=dtype, device=device)
-    k = torch.rand(B, T, H, K, dtype=dtype, device=device)
-    v = torch.rand(B, T, H, V, dtype=dtype, device=device)
-    beta = torch.rand(B, T, H, dtype=torch.float32, device=device).sigmoid().to(dtype)
-    g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32, device=device))
+    q = torch.rand(B, seq_len, H, K, dtype=dtype, device=device)
+    k = torch.rand(B, seq_len, H, K, dtype=dtype, device=device)
+    v = torch.rand(B, seq_len, H, V, dtype=dtype, device=device)
+    beta = (
+        torch.rand(B, seq_len, H, dtype=torch.float32, device=device)
+        .sigmoid()
+        .to(dtype)
+    )
+    g = F.logsigmoid(torch.rand(B, seq_len, H, dtype=torch.float32, device=device))
     g = (g / gate_logit_normalizer).to(dtype)
     if mask_p > 0:
-        g = g * (torch.rand(B, T, H, dtype=dtype, device=device) > mask_p)
+        g = g * (torch.rand(B, seq_len, H, dtype=dtype, device=device) > mask_p)
     if nonzero_h0:
         h0 = torch.randn(B, H, K, V, dtype=dtype, device=device)
     else:
         h0 = torch.zeros(B, H, K, V, dtype=dtype, device=device)
     return q, k, v, g, beta, h0
 
 
-def _make_inputs(dtype=torch.bfloat16, device="cuda"):
-    q = torch.randn(B, T, H, K, dtype=dtype, device=device)
-    k = torch.randn(B, T, H, K, dtype=dtype, device=device)
-    v = torch.randn(B, T, H, V, dtype=dtype, device=device)
-    g = F.logsigmoid(torch.randn(B, T, H, dtype=dtype, device=device))
-    beta = torch.rand(B, T, H, dtype=dtype, device=device).sigmoid()
+def _make_inputs(seq_len=T, dtype=torch.bfloat16, device="cuda"):
+    q = torch.randn(B, seq_len, H, K, dtype=dtype, device=device)
+    k = torch.randn(B, seq_len, H, K, dtype=dtype, device=device)
+    v = torch.randn(B, seq_len, H, V, dtype=dtype, device=device)
+    g = F.logsigmoid(torch.randn(B, seq_len, H, dtype=dtype, device=device))
+    beta = torch.rand(B, seq_len, H, dtype=dtype, device=device).sigmoid()
     initial_state = torch.randn(B, H, K, V, dtype=dtype, device=device)
     return q, k, v, g, beta, initial_state
 
@@ -252,6 +258,69 @@ def test_eager_matches_fla(self):
 
         self.assertLess((o_ours.float() - o_ref.float()).abs().max().item(), 0.01)
 
+    def test_recurrent_t1(self):
+        """T=1 (decode) uses recurrent kernel — verify vs naive reference."""
+        from fla.ops.gated_delta_rule.naive import naive_recurrent_gated_delta_rule
+
+        model = ChunkGatedDeltaModel().eval()
+        for seed, norm, mask_p, nonzero_h0, desc in FLA_TEST_CONFIGS:
+            with self.subTest(desc=desc):
+                inputs = _make_inputs_from_fla(
+                    seed, norm, mask_p, nonzero_h0, seq_len=1
+                )
+                q, k, v, g, beta, h0 = inputs
+
+                with torch.no_grad():
+                    o_ours, s_ours = model(q, k, v, g, beta, h0)
+
+                    o_ref, s_ref = naive_recurrent_gated_delta_rule(
+                        q=F.normalize(q, p=2, dim=-1),
+                        k=F.normalize(k, p=2, dim=-1),
+                        v=v,
+                        beta=beta,
+                        g=g,
+                        initial_state=h0,
+                        output_final_state=True,
+                    )
+
+                self.assertEqual(o_ours.shape, torch.Size([B, 1, H, V]))
+                self.assertEqual(s_ours.shape, torch.Size([B, H, K, V]))
+                o_diff = (o_ours.float() - o_ref.float()).abs().max().item()
+                s_diff = (s_ours.float() - s_ref.float()).abs().max().item()
+                self.assertLess(o_diff, 0.01, f"{desc}: output diff {o_diff}")
+                self.assertLess(s_diff, 0.01, f"{desc}: state diff {s_diff}")
+
+    def test_dispatch_multiple_seq_lengths(self):
+        """Verify correctness across T values hitting both dispatch paths."""
+        from fla.ops.gated_delta_rule.naive import naive_recurrent_gated_delta_rule
+
+        model = ChunkGatedDeltaModel().eval()
+        # T=1 → recurrent, T>1 → chunked; include boundary values
+        for seq_len in [1, 2, 32, 63, 64, 65, 128, 256]:
+            with self.subTest(T=seq_len):
+                inputs = _make_inputs_from_fla(42, 1.0, 0.0, True, seq_len=seq_len)
+                q, k, v, g, beta, h0 = inputs
+
+                with torch.no_grad():
+                    o_ours, s_ours = model(q, k, v, g, beta, h0)
+
+                    o_ref, s_ref = naive_recurrent_gated_delta_rule(
+                        q=F.normalize(q, p=2, dim=-1),
+                        k=F.normalize(k, p=2, dim=-1),
+                        v=v,
+                        beta=beta,
+                        g=g,
+                        initial_state=h0,
+                        output_final_state=True,
+                    )
+
+                self.assertEqual(o_ours.shape, torch.Size([B, seq_len, H, V]))
+                self.assertEqual(s_ours.shape, torch.Size([B, H, K, V]))
+                o_diff = (o_ours.float() - o_ref.float()).abs().max().item()
+                s_diff = (s_ours.float() - s_ref.float()).abs().max().item()
+                self.assertLess(o_diff, 0.02, f"T={seq_len}: output diff {o_diff}")
+                self.assertLess(s_diff, 0.02, f"T={seq_len}: state diff {s_diff}")
+
     def test_export_cuda(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             pte_path = export_chunk_gated_delta(tmpdir)