Use cuDNN for Qwen3.8-27B long-context training

aspec/commands/_recipe_capacity.py

@@ -19,8 +19,14 @@
 import json
 from pathlib import Path
 
+from omegaconf import OmegaConf
+
 
 def audit_long_run(directory):
+    config = OmegaConf.load(Path(directory) / "config.yaml")
+    backend = OmegaConf.select(config, "training.attention_backend")
+    if backend not in {"fa4", "cudnn_flex_attention"}:
+        raise ValueError(f"Unsupported configured attention backend for long-context CP2: {backend!r}")
     evidence = json.loads((Path(directory) / "evidence.json").read_text())
     training = evidence.get("observations", {}).get("training", {})
     scheduler = training.get("scheduler", {})
@@ -32,14 +38,15 @@ def audit_long_run(directory):
         or scheduler.get("rows_dispatched", 0) < 1
         or scheduler.get("packed_tokens") != 261888 * scheduler.get("rows_dispatched", 0)
         or execution.get("context_parallel_size") != 2
-        or execution.get("attention_backend") != "fa4"
+        or execution.get("attention_backend") != backend
     ):
-        raise ValueError("Missing actual 261888-token packed work with FA4 CP2")
+        raise ValueError(f"Missing actual 261888-token packed work with {backend} CP2")
     return {
         "tokens_per_row": 261888,
         "rows": scheduler["rows_dispatched"],
         "packed_tokens": scheduler["packed_tokens"],
         "context_parallel_size": 2,
+        "attention_backend": backend,
     }
 
 

configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml → configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml

@@ -1,9 +1,10 @@
 # Release training starting point: three passes, held-out eval and checkpoints.
 # Hyperparameters are not a convergence/quality claim; see docs/user-guide/recipes.md.
 # Portable execution template. See configs/README.md for required inputs and scope.
-# Adapted from AngelSpec's validated H200 Qwen3.8-27B 256K capacity-unlock
-# lane and expressed through ASpec's canonical schema. One 8-GPU node:
+# Adapted from AngelSpec's Qwen3.8-27B long-context configuration using
+# ASpec's canonical schema. One 8-GPU node:
 # three TP2 Target engines (six GPUs) and one two-rank Direct CP trainer.
+# Uses native cuDNN THD attention for the dense DSpark Draft.
 # Native prefill with optimized direct prefill disabled: one concurrent TP2 request
 # per replica and two source slots for the 261888-token workload.
 
@@ -27,7 +28,7 @@ dataset:
 
 training:
   dflash_packing: true
-  attention_backend: fa4
+  attention_backend: cudnn_flex_attention
   draft_context_parallel_size: 2
   micro_batch_size: 1
   draft_accumulation_steps: 2
@@ -73,6 +74,8 @@ inference:
   inference_batch_size: 1
   sglang:
     tp_size: 2
+    # Allow cold compilation during TP engine startup.
+    init_timeout: 900
     mem_fraction_static: 0.90
     extra_args:
       max_total_tokens: 262144
@@ -112,5 +115,5 @@ debug:
   performance_drain_steps: 1
   timeline_enabled: false
 
-output_dir: outputs/qwen3_8_27b_256k_fa4_cp2_nvlink
-cache_dir: cache/qwen3_8_27b_256k_fa4_cp2_nvlink
+output_dir: outputs/qwen3_8_27b_256k_cudnn_cp2_nvlink
+cache_dir: cache/qwen3_8_27b_256k_cudnn_cp2_nvlink

docs/contributing/validation.md

@@ -89,17 +89,15 @@ bash tools/run_in_env.sh python -m tools.recipes export \
 
 The 27B bounded acceptance check (`accept qwen27b-long`) rejects a training Parquet containing
 any row shorter than 261,888 tokens. It also requires observed 261,888-token sequences and packed rows, positive dispatched work, and actual
-FA4/CP2 execution in both phases. It records device memory every second in `gpu-memory.csv` and
+CP2 execution with the configured attention backend in both phases (cuDNN by default).
+A different observed backend fails the check.
+It records device memory every second in `gpu-memory.csv` and
 per-device peaks/capacities in `acceptance.json`; missing observations fail acceptance. Sampled
 device-wide memory is distinct from an allocator peak and can include other processes, so use
 exclusive allocated GPUs. Synthetic inputs can establish capacity, but cannot establish quality.
 Ordinary `train qwen27b-long` runs may use shorter valid samples; the minimum length is an
 acceptance requirement, not a training input constraint.
 
-A bounded 261,888-token train/resume/export check on eight SM90 GPUs with 143,771 MiB each
-observed peaks of 118,578 MiB on Target and 99,987 MiB on Draft. These are observations for that
-configuration, not minimum-memory requirements or a guarantee for a different workload.
-
 To audit an existing train/resume/export receipt and its saved memory samples:
 
 ```bash

docs/user-guide/recipes.md

@@ -27,11 +27,12 @@ connections on your hardware before training.
 | `qwen27b-long` | Qwen3.8-27B / DSpark | 1 node, 8 GPUs: 3 × TP2 Target + 2D (CP2); 256K-class | Distribute long-context training with CP |
 | `multinode` | Qwen3-8B / DSpark | 2 nodes × 8 GPUs; each 4T + 4D; 4K | Keep features node-local and synchronize gradients across nodes |
 
-The 27B recipe uses a 261,888-token cap, FA4 and CP2. Each TP2 Target replica admits one request,
+The [27B recipe](../../configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml)
+uses a 261,888-token cap, native cuDNN THD attention and CP2. Each TP2 Target replica admits one request,
 with two source slots. `inference.target_prefill_memory_optimized` stays disabled because the direct
 prefill path currently requires TP1. Size feature buffers and backend token capacity for your inputs;
 see [long-context validation](../contributing/validation.md#algorithm-capacity-and-fault-checks) for
-measured capacity and verification commands. Additional examples are under
+capacity checks and verification commands. Additional examples are under
 [configs/examples](../../configs/examples/README.md).
 
 For a short single-node DSpark run with cuDNN attention, CP1 and no packing, use the

tests/test_recipe_capacity.py

@@ -19,8 +19,10 @@
 from tools.validation.recipe_capacity import audit_long_run, summarize_memory
 
 
-@pytest.mark.parametrize("broken", [None, "short", "empty", "cp1", "padding"])
-def test_long_capacity_requires_observed_full_context(tmp_path, broken):
+@pytest.mark.parametrize("backend", ["fa4", "cudnn_flex_attention"])
+@pytest.mark.parametrize("broken", [None, "short", "empty", "cp1", "padding", "backend"])
+def test_long_capacity_requires_observed_full_context(tmp_path, broken, backend):
+    (tmp_path / "config.yaml").write_text(f"training:\n  attention_backend: {backend}\n")
     training = {
         "scheduler": {
             "observed_sequence_min": 261888,
@@ -29,7 +31,7 @@ def test_long_capacity_requires_observed_full_context(tmp_path, broken):
             "rows_dispatched": 4,
             "packed_tokens": 4 * 261888,
         },
-        "execution": {"context_parallel_size": 2, "attention_backend": "fa4"},
+        "execution": {"context_parallel_size": 2, "attention_backend": backend},
     }
     if broken == "short":
         training["scheduler"]["observed_sequence_min"] = 128
@@ -39,12 +41,25 @@ def test_long_capacity_requires_observed_full_context(tmp_path, broken):
         training["execution"]["context_parallel_size"] = 1
     elif broken == "padding":
         training["scheduler"]["packed_tokens"] -= 1
+    elif broken == "backend":
+        training["execution"]["attention_backend"] = (
+            "fa4" if backend == "cudnn_flex_attention" else "cudnn_flex_attention"
+        )
     (tmp_path / "evidence.json").write_text(json.dumps({"observations": {"training": training}}))
     if broken:
         with pytest.raises(ValueError, match="actual"):
             audit_long_run(tmp_path)
     else:
-        assert audit_long_run(tmp_path)["tokens_per_row"] == 261888
+        result = audit_long_run(tmp_path)
+        assert result["tokens_per_row"] == 261888
+        assert result["attention_backend"] == backend
+
+
+def test_long_capacity_rejects_unsupported_configured_backend(tmp_path):
+    (tmp_path / "config.yaml").write_text("training:\n  attention_backend: math\n")
+    (tmp_path / "evidence.json").write_text("{}")
+    with pytest.raises(ValueError, match="configured attention backend"):
+        audit_long_run(tmp_path)
 
 
 def test_sampled_memory_requires_all_devices_and_repeated_samples(tmp_path):

tools/recipe_catalog.json

@@ -115,7 +115,7 @@
   },
   "qwen27b-long": {
     "configs": [
-      "configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml"
+      "configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml"
     ],
     "model": "Qwen3.8-27B",
     "nodes": 1,