Add Qwen3.5 hybrid linear-attention MoE support (122B-A10B-FP8, 397B-A17B-FP8)

README.md

@@ -129,6 +129,8 @@ Validated package versions are `torch 2.10.0+cu130`, `triton 3.6.0`,
 
 - Quickstart and correctness: `Qwen/Qwen3-30B-A3B-Instruct-2507`
 - Published scaling presets: Qwen3-235B-A22B-FP8 and MiniMax-M2.5-FP8
+- Hybrid linear attention: Qwen3.5-122B-A10B-FP8 and Qwen3.5-397B-A17B-FP8,
+  which interleave gated DeltaNet with full attention at 3:1
 
 ### Quickstart
 
@@ -165,6 +167,21 @@ option list):
   --max-new-tokens 64
 ```
 
+Qwen3.5 ships two presets over the same validator. Each pins a checkpoint
+revision, so a result is reproducible:
+
+```bash
+./scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh
+./scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh
+```
+
+Both need a Ray cluster with 8 GPUs (attention TP 4 + MLP TP/EP 4) and a second
+conda environment holding the vLLM reference, named by `VLLM_ENV_NAME` (default
+`vllm-cuda130`). They start from the eager, graph-off, single-microbatch
+baseline; once that passes, set `AFD_SERVER_EXTRA_ARGS=""` to enable both AFD
+decode graphs. Run `./scripts/validate/qwen35_fastafd_alignment.sh --help` for
+the full control list.
+
 ### Large-Scale AFD Experiments
 
 Presets live under `scripts/experiments/afd/` and assume a running Ray GPU cluster,
@@ -207,6 +224,21 @@ Available presets:
 | Qwen3-235B-A22B-FP8 | 16K | 48 requests / attention GPU | `scripts/experiments/afd/qwen3_235b/run_afd_qwen3_235b_a22b_fp8_16k_b48_dynamicnode_mb2_nsys_alignment.sh` |
 | MiniMax-M2.5-FP8 | 8K | 72 requests / attention GPU | `scripts/experiments/afd/minimax_m25/run_afd_minimax_m25_fp8_8k_b72_dynamicnode_mb2_nsys_alignment.sh` |
 | MiniMax-M2.5-FP8 | 16K | 36 requests / attention GPU | `scripts/experiments/afd/minimax_m25/run_afd_minimax_m25_fp8_16k_b36_dynamicnode_mb2_nsys_alignment.sh` |
+| Qwen3.5-122B-A10B-FP8 | 8K | 256 requests / attention GPU | `scripts/experiments/afd/qwen35/run_afd_qwen35_122b_a10b_fp8_8k_b256_2node_mb2_nsys_alignment.sh` |
+| Qwen3.5-397B-A17B-FP8 | 8K | 320 requests / attention GPU | `scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh` |
+
+The Qwen3.5 presets fix their topology rather than sizing it from the cluster:
+four attention workers (DP 4, no tensor parallelism) and four expert workers
+(DP 4, EP 4) — two 4-GPU nodes on the reference setup. They ignore
+`AFD_TOTAL_NODES` and refuse to start on any other worker count, because a
+retained decode window is only free of tensor-parallel all-reduce traffic at
+this layout. `AFD_NUM_MB=2` is the micro-batch ping-pong overlap and is the
+default; `AFD_NUM_MB=1` gives the serialized baseline to compare against.
+
+```bash
+MODEL_PATH=/path/to/Qwen3.5-397B-A17B-FP8 RUN_VLLM_ALIGNMENT=0 NSYS=0 \
+bash scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh
+```
 
 **vLLM cross-check (`RUN_VLLM_ALIGNMENT=1`).** Scores through an expert-parallel
 vLLM baseline (`--all2all-backend deepep_low_latency --moe-backend deep_gemm`),

prompts/qwen35_122b_alignment.txt

@@ -0,0 +1,8 @@
+Explain in one paragraph how recurrent linear attention differs from full attention during autoregressive decoding.
+Write a Python function that returns the first n Fibonacci numbers, then state its time complexity.
+What is 17 multiplied by 23? Show the intermediate arithmetic briefly.
+Summarize why expert parallelism is useful for mixture-of-experts inference.
+Translate "Verification must use measured evidence" into Chinese.
+List three practical checks to perform before launching a distributed GPU job.
+If a cache has 64 free slots and allocates 19, then releases 7, how many free slots remain?
+Describe the purpose of a numerical alignment test in one concise paragraph.

python/minisgl/afd_attention_runtime.py

@@ -253,6 +253,7 @@ def _attach_input_mapping(
         batch.positions_host = positions_host
         batch.positions = positions_host.to(self.state.device, non_blocking=True)
         table_indices = table_indices_host.to(self.state.device, non_blocking=True)
+        batch.state_indices = table_indices
         if plan.phase == "decode":
             batch.afd_req_table_indices_gpu = table_indices
         pos_i64 = batch.positions.to(torch.int64)
@@ -283,6 +284,7 @@ def _build_microbatch_subbatches(
             sb.positions = batch.positions[ts:te]
             sb.input_ids = batch.input_ids[ts:te]
             sb.out_loc = batch.out_loc[ts:te]
+            sb.state_indices = batch.state_indices[ts:te]
             sb.attn_metadata = batch.attn_metadata_mbs[mb]
             subs.append(sb)
         return subs

python/minisgl/afd_attention_worker.py

@@ -17,7 +17,7 @@
 from minisgl.engine.config import EngineConfig
 from minisgl.engine.engine import _adjust_config, _init_tp_communication
 from minisgl.engine.graph import get_free_memory, mem_GB
-from minisgl.kvcache import create_kvcache_pool
+from minisgl.kvcache import create_hybrid_gdn_state_pool, create_kvcache_pool
 from minisgl.utils import div_even, init_logger, nvtx_label, nvtx_range
 
 from .afd_attention_runtime import AfdAttentionRuntime
@@ -75,6 +75,19 @@ def __init__(self, config: EngineConfig):
         self.tp_cpu_group = _init_tp_communication(config, self.dtype)
         free_memory = self._sync_get_memory()[1]
 
+        self.ctx.gdn_state_pool = create_hybrid_gdn_state_pool(
+            config.model_config,
+            max_slots=config.max_running_req + 1,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        if self.ctx.gdn_state_pool is not None:
+            logger.info_rank0(
+                "AFD allocated hybrid GDN state: layers=%d slots=%d bytes_per_slot=%d",
+                len(self.ctx.gdn_state_pool.layer_ids),
+                self.ctx.gdn_state_pool.max_slots,
+                self.ctx.gdn_state_pool.bytes_per_slot,
+            )
         self.num_pages = self._determine_num_pages(free_memory, config)
         self.max_seq_len = min(config.max_seq_len, self.num_pages * config.page_size)
         aligned_max_seq_len = ((self.max_seq_len + 31) // 32) * 32
@@ -127,15 +140,17 @@ def _sync_get_memory(self) -> tuple[int, int]:
         return min_free_memory, max_free_memory
 
     def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> int:
+        new_free_memory = self._sync_get_memory()[1]
         cache_per_page = (
             2
             * config.model_config.head_dim
             * div_even(config.model_config.num_kv_heads, config.tp_info.size, allow_replicate=True)
             * config.page_size
             * self.dtype.itemsize
-            * config.model_config.num_layers
+            * config.model_config.num_full_attention_layers
         )
-        available_memory = int(config.memory_ratio * old_free_memory)
+        state_memory = old_free_memory - new_free_memory
+        available_memory = int(config.memory_ratio * old_free_memory) - state_memory
         memory_cap = available_memory // cache_per_page
         requested = config.num_page_override
         if requested is None:
@@ -311,6 +326,9 @@ def __init__(
         self._input_ids = torch.zeros((self.max_bs,), dtype=torch.int32, device=self.device)
         self._positions = torch.zeros((self.max_bs,), dtype=torch.int32, device=self.device)
         self._out_loc = torch.zeros((self.max_bs,), dtype=torch.int32, device=self.device)
+        self._state_indices = torch.zeros(
+            (self.max_bs,), dtype=torch.int64, device=self.device
+        )
         self._final_out = torch.empty(
             (self.max_bs, self.hidden_size), dtype=torch.bfloat16, device=self.device
         )
@@ -338,6 +356,7 @@ def _build_capture_batch(self, bs: int) -> Batch:
         batch.input_ids = self._input_ids[:bs]
         batch.positions = self._positions[:bs]
         batch.out_loc = self._out_loc[:bs]
+        batch.state_indices = self._state_indices[:bs]
         return batch
 
     def warmup(self, bs: int) -> None:
@@ -465,6 +484,7 @@ def _warmup_mb(self, bs: int) -> None:
             sub.input_ids = self._input_ids[s:e]
             sub.positions = self._positions[s:e]
             sub.out_loc = self._out_loc[s:e]
+            sub.state_indices = self._state_indices[s:e]
             backend.prepare_for_capture(sub)
             subs.append(sub)
         whole.attn_metadata = subs[0].attn_metadata  # embed ignores it
@@ -552,6 +572,7 @@ def _replay_mb(self, batch: Batch, bs: int) -> torch.Tensor:
             self._input_ids[:bs].copy_(batch.input_ids[:bs])
             self._positions[:bs].copy_(batch.positions[:bs])
             self._out_loc[:bs].copy_(batch.out_loc[:bs])
+            self._state_indices[:bs].copy_(batch.state_indices[:bs])
             for mb in range(num_mb):
                 cap_sub = buf.sub_batches[mb]
                 cap_sub.reqs = real_subs[mb].reqs
@@ -576,6 +597,7 @@ def replay(self, batch: Batch, bs: int) -> torch.Tensor:
             self._input_ids[:bs].copy_(batch.input_ids[:bs])
             self._positions[:bs].copy_(batch.positions[:bs])
             self._out_loc[:bs].copy_(batch.out_loc[:bs])
+            self._state_indices[:bs].copy_(batch.state_indices[:bs])
             cap = buf.capture_batch
             cap.reqs = batch.reqs
             cap.padded_reqs = batch.padded_reqs
@@ -878,14 +900,15 @@ def centralized_normal_loop(self) -> None:
                 )
 
     def centralized_overlap_loop(self) -> None:
-        # Fixed two-step-ahead AG: launch two future commands before retiring the
-        # oldest reply handle. This lets cudaGraphLaunch for N+2 enqueue before
-        # the CPU waits on N's D2H copy event. AfdFlushStepCmd is the explicit
-        # pipeline-drain marker at generation end.
+        # Launch a finite future-command window before retiring the oldest reply
+        # handle. With the default pending window 3, cudaGraphLaunch for N+2
+        # enqueues before the CPU waits on N's D2H copy event.
         pending_afd_ag: deque[_AfdAgInFlight] = deque()
 
         def retire_afd_ag(force: bool = False) -> None:
-            while pending_afd_ag and (force or len(pending_afd_ag) > 2):
+            while pending_afd_ag and (
+                force or len(pending_afd_ag) > self.async_retire_depth
+            ):
                 self._process_last_afd_ag(pending_afd_ag.popleft())
 
         while True:
@@ -1038,7 +1061,8 @@ def _run_ag_forward(self, batch, plan: AfdAGStepPlan, ctx) -> torch.Tensor:
             log_line(
                 self.log_path,
                 f"[{self.role} rank={self.tp_rank}] afd_ag_decode_graph:replay "
-                f"step_id={int(plan.step_id)} bs={int(graph_bs)} num_mb={int(plan.num_mb)}",
+                f"step_id={int(plan.step_id)} real={int(plan.real_size)} "
+                f"bs={int(graph_bs)} num_mb={int(plan.num_mb)}",
                 flush=True,
             )
             return self._afd_ag_graph.replay(batch, graph_bs)

python/minisgl/afd_coordinator.py

@@ -131,6 +131,9 @@ def __init__(self, server_args: ServerArgs):
             raise RuntimeError("--afd-max-batched-tokens must be >= 1")
         self.max_seq_len = int(server_args.max_seq_len)
         self.device_comm_num_sms = max(1, int(server_args.afd_device_comm_num_sms))
+        self.async_pending_steps = max(
+            2, int(server_args.afd_async_pending_steps)
+        )
         if str(server_args.cache_type) != "naive":
             raise RuntimeError(
                 "AFD serve uses the centralized scheduler and currently supports only cache_type='naive'"
@@ -243,8 +246,13 @@ def _init_runtime_state(self) -> None:
         self._afd_reply_tokens_by_dp: dict[int, dict[int, list[int]]] = {}
         self._shutdown_requested = False
         self._nsys_runtime_started = False
+        self._nsys_runtime_finished = False
         self._nsys_start_step = max(0, _env_int("MINISGL_RAY_NSYS_START_STEP", 0))
         self._nsys_stop_step = max(0, _env_int("MINISGL_RAY_NSYS_STOP_STEP", 0))
+        self._shutdown_timeout_s = max(
+            5.0,
+            float(_env_int("MINISGL_AFD_SHUTDOWN_TIMEOUT_S", 30)),
+        )
         self._thread: threading.Thread | None = None
         self._failure: str | None = None
 
@@ -474,6 +482,8 @@ def _log_ready_banner(self) -> None:
             f"attention_backend={self.attention_backend} "
             f"max_batched_tokens={self.max_batched_tokens} "
             f"decode_graph_bs={list(self.decode_graph_bs)} "
+            f"disable_overlap={bool(self.server_args.afd_disable_overlap)} "
+            f"async_pending_steps={self.async_pending_steps} "
             f"centralized_scheduler=True",
             flush=True,
         )
@@ -560,11 +570,11 @@ def shutdown(self) -> None:
                 lambda: self._broadcast_cmd_to_workers(AfdStopCmd()),
             )
         if self._thread is not None:
-            self._thread.join(timeout=5.0)
+            self._thread.join(timeout=self._shutdown_timeout_s)
         if self._worker_hot_loop_refs:
             self._shutdown_call(
                 "wait_worker_hot_loops",
-                lambda: ray.get(self._worker_hot_loop_refs, timeout=5.0),
+                self._wait_for_worker_hot_loops,
             )
             self._worker_hot_loop_refs = []
         if self._nsys_runtime_started:
@@ -598,6 +608,12 @@ def shutdown(self) -> None:
             log_line(self.log_path, f"[afd-coordinator] cpu_trace flushed path={trace_path}", flush=True)
         flush_log_lines(self.log_path)
 
+    def _wait_for_worker_hot_loops(self) -> None:
+        ray.get(
+            self._worker_hot_loop_refs,
+            timeout=self._shutdown_timeout_s,
+        )
+
     def _init_centralized_schedulers(self) -> None:
         if not self.attn_workers:
             raise RuntimeError("Cannot initialize centralized scheduler without attention workers")
@@ -1206,7 +1222,7 @@ def _afd_overlap_loop(self) -> None:
         scheds = self._central_schedulers
         attn_tp = int(self.attn_tp_size)
         forward_remains = [_AfdForwardRemain() for _ in range(n_dp)]
-        max_pending = 3
+        max_pending = self.async_pending_steps
         log_line(
             self.log_path,
             f"[afd-overlap] mode=dp_lockstep_fixed_lead n_dp={n_dp} "

python/minisgl/afd_expert_worker.py

@@ -131,7 +131,6 @@ def _run_eg_deepep_pipeline_body(
         [None for _ in range(num_mb)] for _ in range(num_layers)
     ]
     retired: list[tuple[torch.cuda.Event, Any, Any]] = []
-    graph_live_objects: list[Any] = []
 
     def retire(force: bool = False) -> None:
         if not retired:
@@ -187,8 +186,11 @@ def launch_combine(layer: int, mb: int) -> None:
             retired.append((done, disp, expert_out))
             if len(retired) >= 4:
                 retire(force=True)
-        elif torch.cuda.is_current_stream_capturing():
-            graph_live_objects.extend((disp, expert_out))
+        # During graph capture, release each layer's Python objects as soon as
+        # combine has been enqueued. Dispatch storage is created and retired on
+        # the lane stream; expert_out.record_stream above covers its
+        # engine-to-lane handoff. Retaining every layer until capture ends
+        # prevents the graph-private allocator from reusing those buffers.
         dispatches[layer][mb] = None
         expert_outputs[layer][mb] = None
 
@@ -383,10 +385,13 @@ def _fn() -> None:
             fn=_fn,
         )
         self.graphs[total] = graph
+        allocated_gib = torch.cuda.memory_allocated(self.device) / (1024**3)
+        reserved_gib = torch.cuda.memory_reserved(self.device) / (1024**3)
         log_line(
             self.log_path,
             f"[{self.label}] afd_eg_decode_graph warmup:done total={total} "
-            f"mb_bs={mb_bs} schedule=paper_pipeline lanes={len(lane_streams)}",
+            f"mb_bs={mb_bs} schedule=paper_pipeline lanes={len(lane_streams)} "
+            f"allocated_gib={allocated_gib:.2f} reserved_gib={reserved_gib:.2f}",
             flush=True,
         )
 
@@ -599,15 +604,15 @@ def afd_normal_loop(self) -> None:
                 )
 
     def afd_overlap_loop(self) -> None:
-        """Fixed two-step-ahead EG loop.
+        """Finite-window step-ahead EG loop.
 
-        Launch two future commands before retiring the oldest ack handle, matching
-        the AG loop. AfdFlushStepCmd is the explicit pipeline-drain marker.
+        Match the AG worker's configured pending window before retiring the
+        oldest ack handle. AfdFlushStepCmd drains the pipeline.
         """
         pending: deque[_AfdEgInFlight] = deque()
 
         def retire(force: bool = False) -> None:
-            while pending and (force or len(pending) > 2):
+            while pending and (force or len(pending) > self.async_retire_depth):
                 self._process_last_afd_eg(pending.popleft())
 
         while True:

python/minisgl/afd_profiler.py

@@ -15,7 +15,11 @@ def start_nsys_runtime_capture(
     reason: str,
     via_worker_queue: bool = False,
 ) -> None:
-    if not coord.server_args.ray_nsys or coord._nsys_runtime_started:
+    if (
+        not coord.server_args.ray_nsys
+        or coord._nsys_runtime_started
+        or getattr(coord, "_nsys_runtime_finished", False)
+    ):
         return
     if via_worker_queue:
         coord._broadcast_cmd_to_workers(
@@ -58,6 +62,7 @@ def stop_nsys_runtime_capture(
             ]
         )
     coord._nsys_runtime_started = False
+    coord._nsys_runtime_finished = True
     log_line(
         coord.log_path,
         f"[afd-coordinator] nsys profiler:stop reason={reason} "

python/minisgl/afd_support.py

@@ -37,11 +37,15 @@
     "CUDA_HOME",
     "CUDA_PATH",
     "CUDA_NVCC_EXECUTABLE",
+    "CUDA_DEVICE_MAX_CONNECTIONS",
+    "MINISGL_QWEN35_CONV1D_EAGER",
+    "MINISGL_QWEN35_GATE_EAGER",
     "DG_JIT_CACHE_DIR",
     "DG_JIT_NVCC_COMPILER",
     "DG_JIT_USE_NVRTC",
     "TVM_FFI_CACHE_DIR",
     "FLASHINFER_WORKSPACE_BASE",
+    "FLASHINFER_NVCC",
     "TRITON_PTXAS_BLACKWELL_PATH",
     "MAX_JOBS",
     "MINISGL_DEEPEP_BUILD_DIR",