From b4d408216965d965121c98e9fe13a60b3d5ed389 Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:28:45 +0000 Subject: [PATCH 1/6] Add Qwen3.5 hybrid linear-attention MoE support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3.5 interleaves gated DeltaNet with full attention at 3:1 — 122B-A10B is 48 layers (36 GDN, 12 full) and 397B-A17B is 60 (45, 15) — so it needs a KV pool that holds recurrent state and paged KV side by side, which `kvcache/hybrid_pool.py` provides. The GDN decode path lands in `kernel/qwen35_gdn.py`: a causal conv1d that holds its taps in scalar registers through a KERNEL_WIDTH cascade, and a recurrent gated-delta kernel whose grid puts the value block on the fastest- varying axis so consecutive CTAs walk one request's state contiguously. The gated RMSNorm is compiled rather than eager, which is how vLLM gets the same arithmetic as a single inductor kernel. 397B routes 512 experts at top-k 10 against 122B's 256 at 8, so MegaMoE routing, weight loading and the AFD expert plans all had to stop assuming the smaller shape. `models/weight.py` and `models/config.py` carry most of that, including ModelOpt NVFP4 checkpoints, which route to the FlashInfer NVFP4 runner added under `layers/moe/moe_runner/`. The AFD workers gain a configurable step-ahead window. Both loops previously retired the oldest handle once `len(pending) > 2`; that constant decides how far cudaGraphLaunch runs ahead of the CPU wait on the previous step's D2H copy, and the right depth differs by model and topology, so it is now `--afd-async-pending-steps` threaded from ServerArgs through the coordinator. `--afd-force-multi-mb-graph-overlap` is the explicit opt-in for shapes where multi-microbatch graph capture would otherwise auto-serialise. Verified on two GB200 nodes, attention DP4 / TP1 and MLP EP4 / TP1, against a vLLM reference on the same checkpoints. --- python/minisgl/afd_attention_runtime.py | 2 + python/minisgl/afd_attention_worker.py | 42 +- python/minisgl/afd_coordinator.py | 22 +- python/minisgl/afd_expert_worker.py | 21 +- python/minisgl/afd_profiler.py | 7 +- python/minisgl/afd_support.py | 4 + python/minisgl/afd_worker_base.py | 139 +++- python/minisgl/afd_worker_launcher.py | 1 + python/minisgl/core.py | 6 +- python/minisgl/engine/engine.py | 32 +- python/minisgl/engine/graph.py | 4 + python/minisgl/kernel/deepep_moe.py | 9 +- python/minisgl/kernel/moe_topk.py | 4 +- python/minisgl/kernel/nvfp4_moe.py | 542 ++++++++++++++ python/minisgl/kernel/qwen35_gdn.py | 699 ++++++++++++++++++ python/minisgl/kvcache/__init__.py | 10 + python/minisgl/kvcache/hybrid_pool.py | 118 +++ python/minisgl/kvcache/mha_pool.py | 21 +- python/minisgl/layers/__init__.py | 6 +- python/minisgl/layers/attention.py | 4 +- python/minisgl/layers/linear.py | 26 + python/minisgl/layers/moe/layer.py | 61 +- .../minisgl/layers/moe/moe_runner/__init__.py | 4 + python/minisgl/layers/moe/moe_runner/base.py | 4 + .../layers/moe/moe_runner/flashinfer_nvfp4.py | 186 +++++ .../layers/moe/token_dispatcher/deepep.py | 38 +- .../layers/moe/token_dispatcher/standard.py | 16 + python/minisgl/layers/norm.py | 28 + python/minisgl/layers/rotary.py | 24 +- python/minisgl/models/afd.py | 26 +- python/minisgl/models/config.py | 200 ++++- python/minisgl/models/qwen35_moe.py | 391 ++++++++++ python/minisgl/models/qwen35_moe_afd.py | 202 +++++ python/minisgl/models/qwen3_moe_afd.py | 5 +- python/minisgl/models/register.py | 4 + python/minisgl/models/weight.py | 266 ++++++- python/minisgl/moe/deepep_m2n_adapter.py | 11 +- python/minisgl/scheduler/scheduler.py | 6 + python/minisgl/server/args.py | 126 +++- python/minisgl/server/supervisor_ray.py | 3 + python/minisgl/utils/hf.py | 34 +- 41 files changed, 3231 insertions(+), 123 deletions(-) create mode 100644 python/minisgl/kernel/nvfp4_moe.py create mode 100644 python/minisgl/kernel/qwen35_gdn.py create mode 100644 python/minisgl/kvcache/hybrid_pool.py create mode 100644 python/minisgl/layers/moe/moe_runner/flashinfer_nvfp4.py create mode 100644 python/minisgl/models/qwen35_moe.py create mode 100644 python/minisgl/models/qwen35_moe_afd.py diff --git a/python/minisgl/afd_attention_runtime.py b/python/minisgl/afd_attention_runtime.py index 74de798..b84d376 100644 --- a/python/minisgl/afd_attention_runtime.py +++ b/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 diff --git a/python/minisgl/afd_attention_worker.py b/python/minisgl/afd_attention_worker.py index 2217729..895a425 100644 --- a/python/minisgl/afd_attention_worker.py +++ b/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) diff --git a/python/minisgl/afd_coordinator.py b/python/minisgl/afd_coordinator.py index 7254feb..efe957f 100644 --- a/python/minisgl/afd_coordinator.py +++ b/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} " diff --git a/python/minisgl/afd_expert_worker.py b/python/minisgl/afd_expert_worker.py index 66b86a2..ecb8964 100644 --- a/python/minisgl/afd_expert_worker.py +++ b/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: diff --git a/python/minisgl/afd_profiler.py b/python/minisgl/afd_profiler.py index 19da9f0..f0b0c20 100644 --- a/python/minisgl/afd_profiler.py +++ b/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} " diff --git a/python/minisgl/afd_support.py b/python/minisgl/afd_support.py index bf73e88..14acc5e 100644 --- a/python/minisgl/afd_support.py +++ b/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", diff --git a/python/minisgl/afd_worker_base.py b/python/minisgl/afd_worker_base.py index 9c98d80..73eb52d 100644 --- a/python/minisgl/afd_worker_base.py +++ b/python/minisgl/afd_worker_base.py @@ -116,6 +116,7 @@ def __init__( moe_runner_backend: str = "auto", ray_nsys: bool = False, disable_overlap: bool = False, + async_pending_steps: int = 3, afd_num_mb: int = 1, rpc_timeout_ms: int = 300_000, attn_dp_rank: int = 0, @@ -204,6 +205,8 @@ def __init__( self.mlp_ep_size = int(ep_size) self.ep_size = self.mlp_ep_size if role == "mlp" else 1 self.disable_overlap = bool(disable_overlap) + self.async_pending_steps = max(2, int(async_pending_steps)) + self.async_retire_depth = self.async_pending_steps - 1 self.afd_num_mb = max(1, int(afd_num_mb)) self._control_profile_start_step = int( os.environ.get("MINISGL_AFD_CONTROL_PROFILE_START_STEP", "0") or 0 @@ -327,6 +330,10 @@ def __init__( f"attn_tp_size={self.attn_tp_size} mlp_tp_size={self.mlp_tp_size} " f"device_comm_num_sms={self.device_comm_num_sms} " f"afd_num_mb={self.afd_num_mb} " + f"disable_overlap={self.disable_overlap} " + f"async_pending_steps={self.async_pending_steps} " + f"async_retire_depth={self.async_retire_depth} " + f"cuda_device_max_connections={os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS', '')} " f"dual_stream_microbatch={self._dual_stream_microbatch} " f"rpc_timeout_ms={self._recv_timeout_ms} " f"decode_graph_bs={list(self.decode_graph_bs)} " @@ -360,6 +367,49 @@ def _prepare_hot_rpc_runtime(self) -> None: # injection before DeepEP/DeepGEMM JIT compilation starts. self._detach_nsys_injection_from_subprocesses() self._init_afd_runtime() + + # Cold page cache can make EG stage loading minutes slower than AG + # loading. Do not let an early AG rank enter DeepEP graph warmup while + # another union rank is still loading weights: DeepEP's collective + # timeout is shorter than the distributed startup timeout. + if ( + self.enable_decode_graph + and self.decode_graph_bs + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] " + "afd_graph_warmup_barrier:start", + flush=True, + ) + torch.distributed.barrier() + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] " + "afd_graph_warmup_barrier:done", + flush=True, + ) + + # Runtime/model initialization can leave large, releasable blocks in the + # CUDA caching allocator. Graph capture needs physical headroom for its + # private pools, so mirror the ordinary engine graph lifecycle here. + torch.cuda.synchronize(self.device) + allocated = torch.cuda.memory_allocated(self.device) + reserved_before = torch.cuda.memory_reserved(self.device) + torch.cuda.empty_cache() + reserved_after = torch.cuda.memory_reserved(self.device) + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] afd_graph_capture_memory_prepared " + f"allocated_gib={allocated / (1024**3):.2f} " + f"reserved_before_gib={reserved_before / (1024**3):.2f} " + f"reserved_after_gib={reserved_after / (1024**3):.2f} " + f"released_gib={(reserved_before - reserved_after) / (1024**3):.2f}", + flush=True, + ) + # AG and EG graph capture contains matching collectives, so all ranks # self-trigger warmup here before the coordinator starts sending steps. self.warmup_afd_decode_graphs(self.decode_graph_bs) @@ -508,6 +558,9 @@ def _init_afd_runtime(self) -> None: else: if adapter_lanes > 1: os.environ["MINISGL_DEEPEP_PER_BUFFER_COMM_STREAM"] = "1" + deepep_expanded_layout = ( + str(config.afd_moe_runner_backend) != "flashinfer_nvfp4" + ) self._afd_adapters = [ DeepEPM2NAdapter( group=dist.group.WORLD, @@ -517,6 +570,7 @@ def _init_afd_runtime(self) -> None: hidden_size=int(mc.hidden_size), top_k=int(mc.num_experts_per_tok), num_max_dispatch_tokens_per_rank=bucket, + use_expanded_layout=deepep_expanded_layout, ) for _lane in range(adapter_lanes) ] @@ -525,6 +579,7 @@ def _init_afd_runtime(self) -> None: self.log_path, f"[{self.role} rank={self.tp_rank}] afd_deepep_adapters " f"lanes={adapter_lanes} afd_num_mb={int(self.afd_num_mb)} bucket={bucket} " + f"expanded_layout={int(deepep_expanded_layout)} " f"per_buffer_comm_stream={os.environ.get('MINISGL_DEEPEP_PER_BUFFER_COMM_STREAM', '0')}", flush=True, ) @@ -533,6 +588,7 @@ def _init_afd_runtime(self) -> None: # that need FP32 are listed below; everything else follows model dtype. is_minimax_m2 = arch0 == "MiniMaxM2ForCausalLM" is_glm4_moe = arch0 == "Glm4MoeForCausalLM" + is_qwen35 = arch0 == "Qwen3_5MoeForConditionalGeneration" def _minimax_fp32_key(key: str) -> bool: return is_minimax_m2 and ( @@ -546,15 +602,24 @@ def _glm4_fp32_key(key: str) -> bool: or key.endswith(".mlp.gate.e_score_correction_bias") ) + def _qwen35_fp32_key(key: str) -> bool: + return is_qwen35 and ( + key.endswith(".linear_attn.A_log") + or key.endswith("_input_scale") + or key.endswith("_weight_scale_2") + ) + def _cast_full(pairs): tgt = config.dtype return { k: ( v - if v.dtype == torch.float8_e4m3fn - or (v.dtype == torch.float32 and k.endswith("_scale")) + if v.dtype in (torch.uint8, torch.float8_e4m3fn) else v.to(torch.float32) - if _minimax_fp32_key(k) or _glm4_fp32_key(k) + if k.endswith("_scale") + or _minimax_fp32_key(k) + or _glm4_fp32_key(k) + or _qwen35_fp32_key(k) else v.to(tgt) ) for k, v in pairs @@ -569,6 +634,13 @@ def _afd_stage_classes(): ) return Qwen3AfdDenseRouterForCausalLM, Qwen3AfdExpertStage + if arch == "Qwen3_5MoeForConditionalGeneration": + from minisgl.models.qwen35_moe_afd import ( + Qwen3_5AfdDenseRouterForCausalLM, + Qwen3_5AfdExpertStage, + ) + + return Qwen3_5AfdDenseRouterForCausalLM, Qwen3_5AfdExpertStage if arch == "MiniMaxM2ForCausalLM": from minisgl.models.minimax_m2_afd import ( MiniMaxM2AfdDenseRouterForCausalLM, @@ -610,14 +682,14 @@ def _afd_stage_classes(): ) model.load_state_dict(extract_stage_state_dict(model, full)) del full + remote_layers = [ + layer + for layer in getattr(model.model.layers, "op_list", []) + if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts") + ] try: sd = model.state_dict() tensor_bytes = sum(int(t.numel() * t.element_size()) for t in sd.values()) - remote_layers = [ - layer - for layer in getattr(model.model.layers, "op_list", []) - if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts") - ] shapes = [] for idx in (0, len(remote_layers) - 1): if 0 <= idx < len(remote_layers): @@ -645,6 +717,21 @@ def _afd_stage_classes(): f"{type(exc).__name__}: {exc}", flush=True, ) + if str(config.afd_moe_runner_backend) == "flashinfer_nvfp4": + if not remote_layers: + raise RuntimeError("NVFP4 EG prewarm requires at least one expert layer") + experts = remote_layers[0].mlp.experts + started = time.perf_counter() + experts.runner.prewarm(experts, dtype=config.dtype) + torch.cuda.empty_cache() + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] afd_nvfp4_prewarm_done " + f"layer=0 elapsed_s={time.perf_counter() - started:.3f} " + f"cuda_alloc_gib={torch.cuda.memory_allocated() / (1024**3):.2f} " + f"cuda_reserved_gib={torch.cuda.memory_reserved() / (1024**3):.2f}", + flush=True, + ) if self._afd_moe_backend == "megamoe_m2n": # One-time per-layer requant of FP8 expert weights to the mega # kernel's per-32 UE8M0 interleaved format. @@ -897,6 +984,7 @@ def start_hot_rpc_loop( raise finally: self._hot_loop_running = False + self._log_cuda_memory_peak() # Finalize profiling here too: if the coordinator is killed before it # can call worker.shutdown() (observed with attn_dp>1 multi-node # teardown), this is the only place the per-rank cpu_trace + nsys-rep @@ -938,6 +1026,7 @@ def _finalize_profiling(self) -> None: pass def shutdown(self) -> None: + self._log_cuda_memory_peak() self._finalize_profiling() for attr in ( "_recv_from_coordinator", @@ -948,6 +1037,40 @@ def shutdown(self) -> None: self._destroy_afd_adapters() flush_log_lines(self.log_path) + def _log_cuda_memory_peak(self) -> None: + if getattr(self, "_cuda_memory_peak_logged", False): + return + self._cuda_memory_peak_logged = True + try: + current_allocated = int(torch.cuda.memory_allocated(self.device)) + current_reserved = int(torch.cuda.memory_reserved(self.device)) + peak_allocated = int(torch.cuda.max_memory_allocated(self.device)) + peak_reserved = int(torch.cuda.max_memory_reserved(self.device)) + _, total = torch.cuda.mem_get_info(self.device) + gib = 1024**3 + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] afd_cuda_memory " + f"current_allocated_bytes={current_allocated} " + f"current_reserved_bytes={current_reserved} " + f"peak_allocated_bytes={peak_allocated} " + f"peak_reserved_bytes={peak_reserved} " + f"device_total_bytes={int(total)} " + f"current_allocated_gib={current_allocated / gib:.2f} " + f"current_reserved_gib={current_reserved / gib:.2f} " + f"peak_allocated_gib={peak_allocated / gib:.2f} " + f"peak_reserved_gib={peak_reserved / gib:.2f} " + f"device_total_gib={int(total) / gib:.2f}", + flush=True, + ) + except Exception as exc: + log_line( + self.log_path, + f"[{self.role} rank={self.tp_rank}] afd_cuda_memory_log_failed " + f"{type(exc).__name__}: {exc}", + flush=True, + ) + def _destroy_afd_adapters(self) -> None: adapters = list(getattr(self, "_afd_adapters", [])) self._afd_adapters = [] diff --git a/python/minisgl/afd_worker_launcher.py b/python/minisgl/afd_worker_launcher.py index 6d7cc3c..ad54af4 100644 --- a/python/minisgl/afd_worker_launcher.py +++ b/python/minisgl/afd_worker_launcher.py @@ -136,6 +136,7 @@ def _worker_runtime_env(ray_rank: int, node_ip: str) -> dict[str, object]: "device_comm_num_sms": coord.device_comm_num_sms, "ray_nsys": bool(coord.server_args.ray_nsys), "disable_overlap": bool(coord.server_args.afd_disable_overlap), + "async_pending_steps": int(coord.async_pending_steps), "afd_num_mb": coord.afd_num_mb, "rpc_timeout_ms": coord._rpc_timeout_ms, "moe_a2a_backend": coord.server_args.afd_moe_a2a_backend, diff --git a/python/minisgl/core.py b/python/minisgl/core.py index 698fbc7..56aedf2 100644 --- a/python/minisgl/core.py +++ b/python/minisgl/core.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from minisgl.attention import BaseAttnBackend, BaseAttnMetadata - from minisgl.kvcache import BaseCacheHandle, BaseKVCachePool + from minisgl.kvcache import BaseCacheHandle, BaseKVCachePool, HybridGDNStatePool from minisgl.moe import BaseMoeBackend @@ -77,6 +77,9 @@ class Batch: positions: torch.Tensor = field(init=False) positions_host: torch.Tensor = field(init=False) out_loc: torch.Tensor = field(init=False) + # Request table slot per flattened token. This remains a mutable device + # tensor during CUDA graph replay and addresses GDN recurrent state. + state_indices: torch.Tensor = field(init=False) padded_reqs: List[Req] = field(init=False) # this field should be set by attention backend attn_metadata: BaseAttnMetadata = field(init=False) @@ -124,6 +127,7 @@ class Context: attn_backend: BaseAttnBackend = field(init=False) moe_backend: BaseMoeBackend = field(init=False) kv_cache: BaseKVCachePool = field(init=False) + gdn_state_pool: HybridGDNStatePool | None = field(default=None, init=False) mlp_deepep_buffer: object | None = None moe_num_token_non_padded: torch.Tensor | int | None = None moe_deepep_dispatch_max_tokens_per_rank: int | None = None diff --git a/python/minisgl/engine/engine.py b/python/minisgl/engine/engine.py index 4fbca41..dadf9e8 100644 --- a/python/minisgl/engine/engine.py +++ b/python/minisgl/engine/engine.py @@ -9,7 +9,7 @@ from minisgl.attention import create_attention_backend from minisgl.core import Batch, Context, Req, set_global_ctx from minisgl.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info -from minisgl.kvcache import create_kvcache_pool +from minisgl.kvcache import create_hybrid_gdn_state_pool, create_kvcache_pool from minisgl.layers import set_rope_device from minisgl.models import create_model, load_weight from minisgl.moe import create_moe_backend @@ -238,6 +238,19 @@ def __init__(self, config: EngineConfig): self.model.load_state_dict(self._load_weight_state_dict(config)) # ======================= KV cache initialization ======================== + 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( + "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(init_free_memory, config) num_tokens = self.num_pages * config.page_size self.ctx.kv_cache = self.kv_cache = create_kvcache_pool( @@ -313,6 +326,10 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso getattr(config.model_config, "architectures", [""])[0] == "Glm4MoeForCausalLM" ) + is_qwen35 = ( + getattr(config.model_config, "architectures", [""])[0] + == "Qwen3_5MoeForConditionalGeneration" + ) def _minimax_fp32_key(key: str) -> bool: return is_minimax_m2 and ( @@ -326,13 +343,20 @@ def _glm4_fp32_key(key: str) -> bool: or key.endswith(".mlp.gate.e_score_correction_bias") ) + def _qwen35_fp32_key(key: str) -> bool: + return is_qwen35 and ( + key.endswith(".linear_attn.A_log") + or key.endswith("_input_scale") + or key.endswith("_weight_scale_2") + ) + return { k: ( v - if v.dtype == torch.float8_e4m3fn + if v.dtype in (torch.uint8, torch.float8_e4m3fn) or (v.dtype == torch.float32 and k.endswith("_scale")) else v.to(torch.float32) - if _minimax_fp32_key(k) or _glm4_fp32_key(k) + if _minimax_fp32_key(k) or _glm4_fp32_key(k) or _qwen35_fp32_key(k) else v.to(self.dtype) ) for k, v in load_weight(config.model_path, self.device) @@ -346,7 +370,7 @@ def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> in * 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 ) num_pages = config.num_page_override if num_pages is None: diff --git a/python/minisgl/engine/graph.py b/python/minisgl/engine/graph.py index e843174..4c520c5 100644 --- a/python/minisgl/engine/graph.py +++ b/python/minisgl/engine/graph.py @@ -22,6 +22,7 @@ class GraphCaptureBuffer: input_ids: torch.Tensor out_loc: torch.Tensor positions: torch.Tensor + state_indices: torch.Tensor logits: torch.Tensor @classmethod @@ -30,6 +31,7 @@ def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuf input_ids=torch.zeros(bs, dtype=torch.int32, device=device), out_loc=torch.zeros(bs, dtype=torch.int32, device=device), positions=torch.zeros(bs, dtype=torch.int32, device=device), + state_indices=torch.zeros(bs, dtype=torch.int64, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), ) @@ -38,12 +40,14 @@ def set_batch(self, batch: Batch) -> None: batch.input_ids = self.input_ids[_slice] batch.out_loc = self.out_loc[_slice] batch.positions = self.positions[_slice] + batch.state_indices = self.state_indices[_slice] def copy_from(self, batch: Batch) -> None: _slice = slice(batch.padded_size) self.input_ids[_slice] = batch.input_ids self.out_loc[_slice] = batch.out_loc self.positions[_slice] = batch.positions + self.state_indices[_slice] = batch.state_indices def _determine_cuda_graph_bs( diff --git a/python/minisgl/kernel/deepep_moe.py b/python/minisgl/kernel/deepep_moe.py index 0b36bea..826d8db 100644 --- a/python/minisgl/kernel/deepep_moe.py +++ b/python/minisgl/kernel/deepep_moe.py @@ -396,6 +396,7 @@ class _ElasticRuntimeHandle: class DeepEPMoeElasticHandle: handle: Any recv_shape: tuple[int, ...] + use_expanded_layout: bool = True class DeepEPMoeElasticBuffer: @@ -547,6 +548,7 @@ def dispatch( expert_alignment: int = 1, num_max_dispatch_tokens_per_rank: int | None = None, do_cpu_sync: bool = False, + use_expanded_layout: bool = True, ): if topk_ids.dtype != self.topk_idx_dtype: topk_ids = topk_ids.to(self.topk_idx_dtype) @@ -612,7 +614,7 @@ def dispatch( False, True, do_cpu_sync, - True, + bool(use_expanded_layout), True, ) if hidden_states_scale is not None and recv_sf is None: @@ -632,6 +634,7 @@ def dispatch( wrapped = DeepEPMoeElasticHandle( handle=runtime_handle, recv_shape=tuple(int(x) for x in recv_x.shape), + use_expanded_layout=bool(use_expanded_layout), ) return ( recv_x, @@ -650,7 +653,7 @@ def combine( runtime_handle = handle.handle if tuple(expert_output.shape) != tuple(handle.recv_shape): raise RuntimeError( - "DeepEP expanded combine expects expert output to match " + "DeepEP combine expects expert output to match " f"recv_shape={handle.recv_shape}, got={tuple(expert_output.shape)}" ) send = expert_output if expert_output.is_contiguous() else expert_output.contiguous() @@ -674,7 +677,7 @@ def combine( None, False, False, - True, + bool(handle.use_expanded_layout), ) return combined diff --git a/python/minisgl/kernel/moe_topk.py b/python/minisgl/kernel/moe_topk.py index 74810f9..06b393e 100644 --- a/python/minisgl/kernel/moe_topk.py +++ b/python/minisgl/kernel/moe_topk.py @@ -356,9 +356,9 @@ def topk_softmax_group_local( num_tokens, num_experts = gating_output.shape if num_tokens == 0: return - if (num_experts & (num_experts - 1)) != 0 or num_experts > 256: + if (num_experts & (num_experts - 1)) != 0 or num_experts > 512: raise RuntimeError( - "topk_softmax_group_local supports only power-of-two num_experts <= 256, " + "topk_softmax_group_local supports only power-of-two num_experts <= 512, " f"got {num_experts}" ) diff --git a/python/minisgl/kernel/nvfp4_moe.py b/python/minisgl/kernel/nvfp4_moe.py new file mode 100644 index 0000000..1e7a313 --- /dev/null +++ b/python/minisgl/kernel/nvfp4_moe.py @@ -0,0 +1,542 @@ +"""Native ModelOpt NVFP4 routed-expert utilities for Blackwell. + +Checkpoint tensors remain packed FP4 with per-16 FP8-E4M3 block scales. The +only persistent conversion is the byte/scale reorder required by FlashInfer's +TRT-LLM MoE kernel; no BF16 or FP8-weight copy is retained. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, fields +import os +from pathlib import Path +import sys +import warnings + +import torch + + +@contextmanager +def _flashinfer_cuda_target() -> Iterator[Path | None]: + """Temporarily point FlashInfer JIT at Conda's CUDA target tree. + + MoE runner modules are imported by attention and expert workers alike. + Keeping this process-global change scoped prevents a runner import from + redirecting unrelated TVM-FFI extensions such as pynccl to the target-only + CUDA tree. + """ + + candidates: list[Path] = [] + for value in (os.environ.get("CUDA_HOME"), os.environ.get("CUDA_PATH")): + if value: + candidates.append(Path(value)) + prefix = Path(sys.prefix) + candidates.append(prefix) + candidates.extend(sorted((prefix / "targets").glob("*"))) + target: Path | None = None + for root in candidates: + if not (root / "include" / "cuda_bf16.h").is_file(): + continue + target = root + break + if target is None: + yield None + return + + keys = ("CUDA_HOME", "CUDA_PATH", "PATH", "FLASHINFER_NVCC") + previous = {key: os.environ.get(key) for key in keys} + compiler_dirs = tuple( + path + for path in (prefix / "nvvm" / "bin", prefix / "bin", target / "bin") + if path.is_dir() + ) + os.environ["CUDA_HOME"] = str(target) + os.environ["CUDA_PATH"] = str(target) + os.environ["PATH"] = os.pathsep.join( + (*map(str, compiler_dirs), previous.get("PATH") or "") + ) + nvcc = prefix / "bin" / "nvcc" + if nvcc.is_file(): + os.environ["FLASHINFER_NVCC"] = str(nvcc) + try: + yield target + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +@dataclass(frozen=True) +class NvFp4PreparedMoEWeights: + gate_up_proj: torch.Tensor + gate_up_proj_scale: torch.Tensor + down_proj: torch.Tensor + down_proj_scale: torch.Tensor + input_global_scale_inv: torch.Tensor + output1_scale: torch.Tensor + output1_gate_scale: torch.Tensor + output2_scale: torch.Tensor + hidden_size: int + intermediate_size: int + global_num_experts: int + local_expert_offset: int + + @property + def local_num_experts(self) -> int: + return int(self.gate_up_proj.shape[0]) + + def persistent_tensors(self) -> tuple[torch.Tensor, ...]: + return tuple( + value + for field in fields(self) + if isinstance((value := getattr(self, field.name)), torch.Tensor) + ) + + +def pack_topk_ids_weights( + topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """Pack expert id and BF16 route weight in FlashInfer/vLLM's int32 format.""" + + if topk_ids.shape != topk_weights.shape: + raise ValueError( + f"top-k id/weight shape mismatch: {topk_ids.shape} != {topk_weights.shape}" + ) + return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view( + torch.int16 + ) + + +def _unpack_fp4_e2m1(packed: torch.Tensor) -> torch.Tensor: + if packed.dtype not in (torch.uint8, torch.int8, torch.float4_e2m1fn_x2): + raise TypeError(f"packed NVFP4 tensor must be byte/float4, got {packed.dtype}") + byte = packed.view(torch.uint8) + code = torch.empty( + (*byte.shape[:-1], byte.shape[-1] * 2), + dtype=torch.uint8, + device=byte.device, + ) + code[..., 0::2] = byte & 0x0F + code[..., 1::2] = byte >> 4 + values = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=byte.device, + ) + magnitude = values[(code & 0x07).long()] + return torch.where((code & 0x08) != 0, -magnitude, magnitude) + + +def _expand_global_weight_scale( + global_scale: torch.Tensor, + *, + weight_shape: tuple[int, ...], +) -> torch.Tensor: + scale = global_scale.float() + if scale.numel() == 1: + return scale + if len(weight_shape) != 3: + raise ValueError( + "non-scalar NVFP4 global scale requires batched expert weights, " + f"got weight={weight_shape} scale={tuple(scale.shape)}" + ) + experts, output_rows, _ = weight_shape + if scale.shape == (experts,): + return scale.view(experts, 1, 1) + if scale.shape == (experts, 2): + if output_rows % 2: + raise ValueError("merged gate/up output rows must be even") + return scale.repeat_interleave(output_rows // 2, dim=1).unsqueeze(-1) + raise ValueError( + f"unsupported NVFP4 global scale shape {tuple(scale.shape)} for {weight_shape}" + ) + + +def dequantize_nvfp4_weight( + packed: torch.Tensor, + block_scale: torch.Tensor, + global_scale: torch.Tensor, + *, + group_size: int = 16, +) -> torch.Tensor: + """Independent FP32 dequantization oracle for ModelOpt NVFP4 weights.""" + + values = _unpack_fp4_e2m1(packed) + if block_scale.dtype != torch.float8_e4m3fn: + raise TypeError( + f"NVFP4 block scales must be FP8-E4M3, got {block_scale.dtype}" + ) + if block_scale.shape[:-1] != packed.shape[:-1]: + raise ValueError( + f"NVFP4 scale prefix mismatch: {block_scale.shape} vs {packed.shape}" + ) + logical_k = packed.shape[-1] * 2 + if block_scale.shape[-1] * group_size != logical_k: + raise ValueError( + f"NVFP4 group-{group_size} scale K mismatch: packed logical K={logical_k}, " + f"scale={block_scale.shape[-1]}" + ) + per_value_scale = torch.repeat_interleave( + block_scale.float(), group_size, dim=-1 + ) + tensor_scale = _expand_global_weight_scale( + global_scale, weight_shape=tuple(values.shape) + ) + return values * per_value_scale * tensor_scale + + +def _reorder_gate_up( + weight: torch.Tensor, scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert checkpoint [gate, up] rows to TRT-LLM's [up, gate] rows.""" + + if weight.shape[-2] % 2: + raise ValueError(f"merged gate/up rows must be even, got {weight.shape}") + half = weight.shape[-2] // 2 + gate, up = weight.split(half, dim=-2) + gate_scale, up_scale = scale.split(half, dim=-2) + return ( + torch.cat((up, gate), dim=-2).contiguous(), + torch.cat((up_scale, gate_scale), dim=-2).contiguous(), + ) + + +def _prepare_static_weights_for_trtllm_configured( + gate_up_proj: torch.Tensor, + down_proj: torch.Tensor, + gate_up_scale: torch.Tensor, + down_scale: torch.Tensor, + *, + hidden_size: int, + intermediate_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Reorder native tensors exactly as FlashInfer TRT-LLM NVFP4 MoE expects.""" + + from flashinfer import nvfp4_block_scale_interleave + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + num_experts = int(gate_up_proj.shape[0]) + epilogue_tile_m = 128 + cache: dict[torch.Size, torch.Tensor] = {} + w1_out: list[torch.Tensor] = [] + s1_out: list[torch.Tensor] = [] + w2_out: list[torch.Tensor] = [] + s2_out: list[torch.Tensor] = [] + + w1 = gate_up_proj.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // 2 + ) + s1 = gate_up_scale.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // 16 + ) + w2 = down_proj.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 2 + ) + s2 = down_scale.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 16 + ) + + for expert in range(num_experts): + w1_idx = _maybe_get_cached_w3_w1_permute_indices( + cache, + w1[expert].view(torch.uint8), + epilogue_tile_m, + is_gated_act_gemm=True, + ) + w1_out.append( + w1[expert].view(torch.uint8)[w1_idx.to(w1.device)].contiguous() + ) + s1_idx = _maybe_get_cached_w3_w1_permute_indices( + cache, + s1[expert].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + is_gated_act_gemm=True, + ) + s1_out.append( + nvfp4_block_scale_interleave( + s1[expert].view(torch.uint8)[s1_idx.to(s1.device)].contiguous() + ) + ) + + w2_idx = get_w2_permute_indices_with_cache( + cache, w2[expert].view(torch.uint8), epilogue_tile_m + ) + w2_out.append( + w2[expert].view(torch.uint8)[w2_idx.to(w2.device)].contiguous() + ) + s2_idx = get_w2_permute_indices_with_cache( + cache, + s2[expert].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + s2_out.append( + nvfp4_block_scale_interleave( + s2[expert].view(torch.uint8)[s2_idx.to(s2.device)].contiguous() + ) + ) + + prepared_s1 = ( + torch.stack(s1_out) + .view(torch.float8_e4m3fn) + .reshape(num_experts, 2 * intermediate_size, hidden_size // 16) + ) + prepared_s2 = ( + torch.stack(s2_out) + .view(torch.float8_e4m3fn) + .reshape(num_experts, hidden_size, intermediate_size // 16) + ) + return torch.stack(w1_out), prepared_s1, torch.stack(w2_out), prepared_s2 + + +def _prepare_static_weights_for_trtllm( + gate_up_proj: torch.Tensor, + down_proj: torch.Tensor, + gate_up_scale: torch.Tensor, + down_scale: torch.Tensor, + *, + hidden_size: int, + intermediate_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + with _flashinfer_cuda_target(): + return _prepare_static_weights_for_trtllm_configured( + gate_up_proj, + down_proj, + gate_up_scale, + down_scale, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + ) + + +def _validate_checkpoint_tensors( + gate_up_proj: torch.Tensor, + gate_up_scale: torch.Tensor, + gate_up_scale_2: torch.Tensor, + down_proj: torch.Tensor, + down_scale: torch.Tensor, + down_scale_2: torch.Tensor, +) -> tuple[int, int, int]: + if gate_up_proj.dtype != torch.uint8 or down_proj.dtype != torch.uint8: + raise TypeError( + "ModelOpt NVFP4 expert weights must remain packed uint8; got " + f"gate_up={gate_up_proj.dtype} down={down_proj.dtype}" + ) + if ( + gate_up_scale.dtype != torch.float8_e4m3fn + or down_scale.dtype != torch.float8_e4m3fn + ): + raise TypeError("ModelOpt NVFP4 per-16 scales must be FP8-E4M3") + if gate_up_proj.dim() != 3 or down_proj.dim() != 3: + raise ValueError("NVFP4 expert weights must have [E, N, packed-K] layout") + experts, twice_intermediate, packed_hidden = gate_up_proj.shape + hidden = packed_hidden * 2 + intermediate = twice_intermediate // 2 + expected = { + "gate_up_scale": (experts, twice_intermediate, hidden // 16), + "down_proj": (experts, hidden, intermediate // 2), + "down_scale": (experts, hidden, intermediate // 16), + "gate_up_scale_2": (experts, 2), + "down_scale_2": (experts,), + } + actual = { + "gate_up_scale": tuple(gate_up_scale.shape), + "down_proj": tuple(down_proj.shape), + "down_scale": tuple(down_scale.shape), + "gate_up_scale_2": tuple(gate_up_scale_2.shape), + "down_scale_2": tuple(down_scale_2.shape), + } + bad = {name: (actual[name], shape) for name, shape in expected.items() if actual[name] != shape} + if bad: + raise ValueError(f"invalid ModelOpt NVFP4 expert tensor shapes: {bad}") + if hidden % 512 != 0: + raise ValueError(f"FlashInfer TRT-LLM NVFP4 MoE requires hidden % 512 == 0, got {hidden}") + if intermediate % 16 != 0: + raise ValueError(f"NVFP4 intermediate size must be divisible by 16, got {intermediate}") + return experts, hidden, intermediate + + +def prepare_nvfp4_moe_weights( + gate_up_proj: torch.Tensor, + gate_up_proj_scale: torch.Tensor, + gate_up_proj_weight_scale_2: torch.Tensor, + gate_up_proj_input_scale: torch.Tensor, + down_proj: torch.Tensor, + down_proj_scale: torch.Tensor, + down_proj_weight_scale_2: torch.Tensor, + down_proj_input_scale: torch.Tensor, + *, + global_num_experts: int | None = None, + local_expert_offset: int = 0, +) -> NvFp4PreparedMoEWeights: + """Prepare one local expert shard without changing its FP4 values.""" + + experts, hidden, intermediate = _validate_checkpoint_tensors( + gate_up_proj, + gate_up_proj_scale, + gate_up_proj_weight_scale_2, + down_proj, + down_proj_scale, + down_proj_weight_scale_2, + ) + if gate_up_proj.device.type != "cuda": + raise RuntimeError("native NVFP4 weight preparation requires CUDA") + if not torch.allclose( + gate_up_proj_weight_scale_2[:, 0], gate_up_proj_weight_scale_2[:, 1] + ): + warnings.warn( + "gate_proj and up_proj weight_scale_2 differ; matching vLLM by using gate_proj", + RuntimeWarning, + stacklevel=2, + ) + + input1 = gate_up_proj_input_scale.float().max() + input2 = down_proj_input_scale.float().max() + if not bool((input1 > 0).item()) or not bool((input2 > 0).item()): + raise ValueError("ModelOpt NVFP4 input global scales must be positive") + w1_global = gate_up_proj_weight_scale_2[:, 0].float().contiguous() + w2_global = down_proj_weight_scale_2.float().contiguous() + + reordered_w1, reordered_s1 = _reorder_gate_up( + gate_up_proj, gate_up_proj_scale + ) + prepared_w1, prepared_s1, prepared_w2, prepared_s2 = ( + _prepare_static_weights_for_trtllm( + reordered_w1, + down_proj, + reordered_s1, + down_proj_scale, + hidden_size=hidden, + intermediate_size=intermediate, + ) + ) + output1_gate_scale = (w1_global * input1).contiguous() + output1_scale = (output1_gate_scale / input2).contiguous() + output2_scale = (w2_global * input2).contiguous() + return NvFp4PreparedMoEWeights( + gate_up_proj=prepared_w1, + gate_up_proj_scale=prepared_s1, + down_proj=prepared_w2, + down_proj_scale=prepared_s2, + input_global_scale_inv=(1.0 / input1).reshape(1).contiguous(), + output1_scale=output1_scale, + output1_gate_scale=output1_gate_scale, + output2_scale=output2_scale, + hidden_size=hidden, + intermediate_size=intermediate, + global_num_experts=int(global_num_experts or experts), + local_expert_offset=int(local_expert_offset), + ) + + +def _run_nvfp4_moe_configured( + prepared: NvFp4PreparedMoEWeights, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Run pre-routed native NVFP4 experts through FlashInfer TRT-LLM MoE.""" + + if hidden_states.device.type != "cuda": + raise RuntimeError("native NVFP4 MoE requires CUDA") + if hidden_states.dtype not in (torch.bfloat16, torch.float16): + raise TypeError(f"NVFP4 MoE input must be BF16/FP16, got {hidden_states.dtype}") + if hidden_states.shape[-1] != prepared.hidden_size: + raise ValueError( + f"hidden size mismatch: {hidden_states.shape[-1]} != {prepared.hidden_size}" + ) + + import flashinfer + from flashinfer.fp4_quantization import fp4_quantize + + # The TRT-LLM MoE launcher consumes a true row-major [M, K/16] activation + # scale tensor. FlashInfer's higher-level nvfp4_quantize wrapper always + # emits a padded 8x4/128x4 swizzle, including for its `layout_linear` enum. + # Use the lower-level unswizzled path, matching vLLM's + # moe_kernel_quantize_input(..., is_fp4_scale_swizzled=False). + hidden_q, hidden_scale = fp4_quantize( + hidden_states.contiguous(), + prepared.input_global_scale_inv, + sf_vec_size=16, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + is_sf_8x4_layout=False, + ) + kernel_output = torch.empty_like(hidden_states) + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( + topk_ids=pack_topk_ids_weights(topk_ids, topk_weights).contiguous(), + routing_bias=None, + hidden_states=hidden_q, + hidden_states_scale=hidden_scale.view(torch.float8_e4m3fn), + gemm1_weights=prepared.gate_up_proj, + gemm1_weights_scale=prepared.gate_up_proj_scale, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=prepared.down_proj, + gemm2_weights_scale=prepared.down_proj_scale, + gemm2_bias=None, + output1_scale_scalar=prepared.output1_scale, + output1_scale_gate_scalar=prepared.output1_gate_scale, + output2_scale_scalar=prepared.output2_scale, + num_experts=prepared.global_num_experts, + top_k=int(topk_ids.shape[1]), + n_group=0, + topk_group=0, + intermediate_size=prepared.intermediate_size, + local_expert_offset=prepared.local_expert_offset, + local_num_experts=prepared.local_num_experts, + routed_scaling_factor=None, + routing_method_type=1, + do_finalize=True, + activation_type=3, + output=kernel_output, + ) + if output is None: + return kernel_output + if output.shape != hidden_states.shape: + raise ValueError( + f"output shape mismatch: {output.shape} != {hidden_states.shape}" + ) + output.copy_(kernel_output) + return output + + +def run_nvfp4_moe( + prepared: NvFp4PreparedMoEWeights, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Run native NVFP4 MoE while containing FlashInfer's CUDA JIT config.""" + + with _flashinfer_cuda_target(): + return _run_nvfp4_moe_configured( + prepared, + hidden_states, + topk_ids, + topk_weights, + output=output, + ) + + +__all__ = [ + "NvFp4PreparedMoEWeights", + "dequantize_nvfp4_weight", + "pack_topk_ids_weights", + "prepare_nvfp4_moe_weights", + "run_nvfp4_moe", +] diff --git a/python/minisgl/kernel/qwen35_gdn.py b/python/minisgl/kernel/qwen35_gdn.py new file mode 100644 index 0000000..94c7bf2 --- /dev/null +++ b/python/minisgl/kernel/qwen35_gdn.py @@ -0,0 +1,699 @@ +from __future__ import annotations + +import math +import os +from typing import Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + + +def qwen35_l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + """Qwen3.5 Q/K L2 normalization, evaluated in FP32.""" + + x_f32 = x.float() + return (x_f32 * torch.rsqrt(x_f32.square().sum(dim=-1, keepdim=True) + eps)).to( + x.dtype + ) + + +def qwen35_gdn_gates( + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return log forget gate ``g`` and delta update gate ``beta`` in FP32.""" + + g = -A_log.float().exp() * F.softplus(a.float() + dt_bias.float()) + beta = torch.sigmoid(b.float()) + return g, beta + + +def _expand_qk_heads( + q: torch.Tensor, k: torch.Tensor, num_value_heads: int +) -> Tuple[torch.Tensor, torch.Tensor]: + num_key_heads = q.shape[-2] + if num_key_heads == num_value_heads: + return q, k + if num_value_heads % num_key_heads != 0: + raise ValueError( + "GDN value heads must be divisible by key heads: " + f"{num_value_heads} % {num_key_heads} != 0" + ) + repeats = num_value_heads // num_key_heads + return ( + q.repeat_interleave(repeats, dim=-2), + k.repeat_interleave(repeats, dim=-2), + ) + + +def qwen35_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + *, + scale: float | None = None, + normalize_qk: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Portable GDN oracle/fallback for one sequence. + + Inputs use ``[T, H, D]`` and state uses the FlashInfer-compatible + k-last layout ``[H_v, D_v, D_k]``. + """ + + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("q, k, and v must have shape [tokens, heads, head_dim]") + if initial_state.shape != ( + v.shape[1], + v.shape[2], + k.shape[2], + ): + raise ValueError( + "invalid recurrent state shape: " + f"got {tuple(initial_state.shape)}, expected " + f"{(v.shape[1], v.shape[2], k.shape[2])}" + ) + + if normalize_qk: + q, k = qwen35_l2norm(q), qwen35_l2norm(k) + q, k = _expand_qk_heads(q, k, v.shape[1]) + scale = float(scale if scale is not None else q.shape[-1] ** -0.5) + state = initial_state.float().clone() + outputs = [] + for token_idx in range(q.shape[0]): + state.mul_(torch.exp(g[token_idx].float())[:, None, None]) + q_t = q[token_idx].float() * scale + k_t = k[token_idx].float() + v_t = v[token_idx].float() + memory = (state * k_t[:, None, :]).sum(dim=-1) + delta = (v_t - memory) * beta[token_idx].float()[:, None] + state.add_(delta[:, :, None] * k_t[:, None, :]) + outputs.append((state * q_t[:, None, :]).sum(dim=-1)) + output = torch.stack(outputs, dim=0).to(q.dtype) + return output, state + + +def qwen35_recurrent_gated_delta_rule_packed( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Portable packed-sequence wrapper used by CPU tests and fallback paths.""" + + boundaries = [int(x) for x in cu_seqlens.to("cpu").tolist()] + if len(boundaries) != initial_state.shape[0] + 1: + raise ValueError("cu_seqlens and initial_state batch size disagree") + outputs = [] + final_states = [] + for seq_idx, (start, end) in enumerate(zip(boundaries, boundaries[1:])): + output, final_state = qwen35_recurrent_gated_delta_rule( + q[start:end], + k[start:end], + v[start:end], + g[start:end], + beta[start:end], + initial_state[seq_idx], + ) + outputs.append(output) + final_states.append(final_state) + return torch.cat(outputs, dim=0), torch.stack(final_states, dim=0) + + +def qwen35_gdn_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """GDN prefill using FlashInfer's public SM90 API when available.""" + + if not q.is_cuda: + return qwen35_recurrent_gated_delta_rule_packed( + q, k, v, g, beta, initial_state, cu_seqlens + ) + + device_major = torch.cuda.get_device_capability(q.device)[0] + if device_major == 9 and q.dtype in (torch.float16, torch.bfloat16): + from flashinfer import chunk_gated_delta_rule + + q_norm = qwen35_l2norm(q).contiguous() + k_norm = qwen35_l2norm(k).contiguous() + return chunk_gated_delta_rule( + q=q_norm, + k=k_norm, + v=v.contiguous(), + g=torch.exp(g.float()).contiguous(), + beta=beta.float().contiguous(), + scale=q.shape[-1] ** -0.5, + initial_state=initial_state.float().contiguous(), + output_final_state=True, + cu_seqlens=cu_seqlens.to(device=q.device, dtype=torch.int64), + use_qk_l2norm_in_kernel=False, + ) + return _qwen35_gdn_prefill_triton(q, k, v, g, beta, initial_state, cu_seqlens) + + +def qwen35_causal_conv1d_prefill( + x: torch.Tensor, + weight: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + first_positions: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Packed depthwise causal convolution and final short-conv state.""" + + width = weight.shape[-1] + if initial_state.shape[-1] != width - 1: + raise ValueError("convolution state width must equal kernel width - 1") + boundaries = [int(v) for v in cu_seqlens.to("cpu").tolist()] + outputs = [] + final_states = [] + for seq_idx, (start, end) in enumerate(zip(boundaries, boundaries[1:])): + reset = (first_positions[seq_idx] == 0).to(initial_state.dtype) + prefix = initial_state[seq_idx] * (1 - reset) + sequence = x[start:end].transpose(0, 1) + window = torch.cat((prefix, sequence), dim=-1) + output = F.conv1d( + window.unsqueeze(0), + weight, + bias=None, + groups=x.shape[-1], + ).squeeze(0).transpose(0, 1) + outputs.append(F.silu(output)) + final_states.append(window[:, -(width - 1) :]) + return torch.cat(outputs, dim=0), torch.stack(final_states, dim=0) + + +@triton.jit +def _qwen35_causal_conv1d_decode_kernel( + x_ptr, + w_ptr, + state_ptr, + idx_ptr, + pos_ptr, + out_ptr, + num_channels, + x_row_stride, + state_slot_stride, + state_channel_stride, + state_tap_stride, + w_channel_stride, + w_tap_stride, + out_row_stride, + KERNEL_WIDTH: tl.constexpr, + BLOCK_C: tl.constexpr, +): + """One fused decode step of the Qwen3.5 depthwise causal convolution. + + Replaces index_select + mask + cat + weighted-sum + index_copy_ + SiLU, which + ran as roughly seven ATen launches per linear-attention layer. + + The retained taps are held in scalar registers instead of a [BLOCK_C, taps] + tile, mirroring vLLM's `_causal_conv1d_update_kernel`. That removes a second + read of the state (the shifted window is already in registers) and the + power-of-two tap padding, which together cost ~1.6x the reference's runtime. + """ + req = tl.program_id(0) + offs_c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + mask_c = offs_c < num_channels + + slot = tl.load(idx_ptr + req) + # A request at position 0 starts a fresh sequence, so its retained state is + # discarded rather than convolved. + keep = tl.where(tl.load(pos_ptr + req) != 0, 1.0, 0.0) + + state_base = state_ptr + slot * state_slot_stride + offs_c * state_channel_stride + w_base = w_ptr + offs_c * w_channel_stride + xv = tl.load(x_ptr + req * x_row_stride + offs_c, mask=mask_c, other=0.0).to(tl.float32) + + # window = [state[0..W-2], x]; accumulate against the matching weight taps. + acc = tl.zeros([BLOCK_C], dtype=tl.float32) + if KERNEL_WIDTH >= 2: + tap0 = tl.load(state_base, mask=mask_c, other=0.0).to(tl.float32) * keep + acc += tap0 * tl.load(w_base, mask=mask_c, other=0.0).to(tl.float32) + if KERNEL_WIDTH >= 3: + tap1 = tl.load( + state_base + state_tap_stride, mask=mask_c, other=0.0 + ).to(tl.float32) * keep + acc += tap1 * tl.load( + w_base + w_tap_stride, mask=mask_c, other=0.0 + ).to(tl.float32) + if KERNEL_WIDTH >= 4: + tap2 = tl.load( + state_base + 2 * state_tap_stride, mask=mask_c, other=0.0 + ).to(tl.float32) * keep + acc += tap2 * tl.load( + w_base + 2 * w_tap_stride, mask=mask_c, other=0.0 + ).to(tl.float32) + acc += xv * tl.load( + w_base + (KERNEL_WIDTH - 1) * w_tap_stride, mask=mask_c, other=0.0 + ).to(tl.float32) + + # Shift the retained window by one: new_state[j] = window[j+1]. Every source + # tap is already in a register, so the state is never re-read. + state_ty = state_ptr.dtype.element_ty + if KERNEL_WIDTH >= 3: + tl.store(state_base, tap1.to(state_ty), mask=mask_c) + if KERNEL_WIDTH >= 4: + tl.store(state_base + state_tap_stride, tap2.to(state_ty), mask=mask_c) + tl.store( + state_base + (KERNEL_WIDTH - 2) * state_tap_stride, + xv.to(state_ty), + mask=mask_c, + ) + + out = acc * tl.sigmoid(acc) + tl.store( + out_ptr + req * out_row_stride + offs_c, + out.to(out_ptr.dtype.element_ty), + mask=mask_c, + ) + + +def qwen35_causal_conv1d_decode_eager( + x: torch.Tensor, + weight: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """Reference implementation retained for parity testing and CPU fallback.""" + + state = state_pool.index_select(0, state_indices) + keep = (positions != 0).to(state.dtype)[:, None, None] + state = state * keep + window = torch.cat((state, x.unsqueeze(-1)), dim=-1) + output = (window * weight.squeeze(1).unsqueeze(0)).sum(dim=-1) + state_pool.index_copy_(0, state_indices, window[:, :, 1:]) + return F.silu(output) + + +def qwen35_causal_conv1d_decode( + x: torch.Tensor, + weight: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """Graph-safe single-token causal convolution with in-place slot updates.""" + + # MINISGL_QWEN35_CONV1D_EAGER=1 selects the pre-fusion path so an A/B can be + # interleaved without editing or checking out source between runs. + if not x.is_cuda or os.environ.get("MINISGL_QWEN35_CONV1D_EAGER") == "1": + return qwen35_causal_conv1d_decode_eager( + x, weight, state_pool, state_indices, positions + ) + + num_requests, num_channels = x.shape + kernel_width = weight.shape[-1] + if not 2 <= kernel_width <= 4: + raise ValueError( + f"qwen35_causal_conv1d_decode supports kernel_width 2..4, got " + f"{kernel_width} from weight shape {tuple(weight.shape)}; the decode " + f"kernel unrolls the retained taps into registers. Extend the " + f"KERNEL_WIDTH cascade in _qwen35_causal_conv1d_decode_kernel to " + f"support a wider convolution." + ) + block_c = 128 + out = torch.empty_like(x) + grid = (num_requests, triton.cdiv(num_channels, block_c)) + _qwen35_causal_conv1d_decode_kernel[grid]( + x, + weight, + state_pool, + state_indices, + positions, + out, + num_channels, + x.stride(0), + state_pool.stride(0), + state_pool.stride(1), + state_pool.stride(2), + weight.stride(0), + weight.stride(-1), + out.stride(0), + KERNEL_WIDTH=kernel_width, + BLOCK_C=block_c, + num_warps=4, + ) + return out + + +def _qwen35_rmsnorm_gated( + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + original_dtype = x.dtype + x_f32 = x.float() + # z was upcast twice here, which cost an extra full-size copy per layer. + z_f32 = z.float() + normalized = x_f32 * torch.rsqrt(x_f32.square().mean(dim=-1, keepdim=True) + eps) + normalized = normalized * weight.float() + gated = normalized * (z_f32 * torch.sigmoid(z_f32)) + return gated.to(original_dtype) + + +# Eager, this is ~13 ATen launches per linear-attention layer -- 36 of them per +# decode step, each materializing an fp32 intermediate the size of the value +# tensor. vLLM gets the same arithmetic as a single inductor kernel because its +# decoder carries `@support_torch_compile`; minisgl already uses this same local +# pattern for the MiniMax QK norm in `layers/norm.py`. +_qwen35_rmsnorm_gated_compiled = torch.compile( + _qwen35_rmsnorm_gated, dynamic=True, fullgraph=True +) + + +def qwen35_rmsnorm_gated( + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Per-value-head RMSNorm followed by the Qwen3.5 SiLU output gate.""" + + # MINISGL_QWEN35_GATE_EAGER=1 selects the pre-compile path so an A/B can be + # interleaved without editing or checking out source between runs, matching + # MINISGL_QWEN35_CONV1D_EAGER above. Needed because the compiled reduction is + # not bitwise-identical, so any correctness difference has to be attributable. + if not x.is_cuda or os.environ.get("MINISGL_QWEN35_GATE_EAGER") == "1": + return _qwen35_rmsnorm_gated(x, z, weight, eps) + return _qwen35_rmsnorm_gated_compiled(x, z, weight, eps) + + +@triton.jit +def _qwen35_gdn_prefill_kernel( + q_ptr, + k_ptr, + v_ptr, + g_ptr, + beta_ptr, + initial_state_ptr, + cu_seqlens_ptr, + output_ptr, + final_state_ptr, + H_Q: tl.constexpr, + H_V: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + sequence_idx = tl.program_id(0) + value_head = tl.program_id(1) + value_block = tl.program_id(2) + key_head = value_head // (H_V // H_Q) + start = tl.load(cu_seqlens_ptr + sequence_idx) + end = tl.load(cu_seqlens_ptr + sequence_idx + 1) + + key_offsets = tl.arange(0, BK) + value_offsets = value_block * BV + tl.arange(0, BV) + key_mask = key_offsets < K + value_mask = value_offsets < V + state_offsets = ( + ((sequence_idx * H_V + value_head) * V + value_offsets[:, None]) * K + + key_offsets[None, :] + ) + state_mask = value_mask[:, None] & key_mask[None, :] + state = tl.load( + initial_state_ptr + state_offsets, mask=state_mask, other=0.0 + ).to(tl.float32) + + token_idx = start + while token_idx < end: + q = tl.load( + q_ptr + (token_idx * H_Q + key_head) * K + key_offsets, + mask=key_mask, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_ptr + (token_idx * H_Q + key_head) * K + key_offsets, + mask=key_mask, + other=0.0, + ).to(tl.float32) + q *= tl.rsqrt(tl.sum(q * q, axis=0) + 1e-6) * (K**-0.5) + k *= tl.rsqrt(tl.sum(k * k, axis=0) + 1e-6) + head_offset = token_idx * H_V + value_head + state *= tl.exp(tl.load(g_ptr + head_offset).to(tl.float32)) + v = tl.load( + v_ptr + head_offset * V + value_offsets, + mask=value_mask, + other=0.0, + ).to(tl.float32) + memory = tl.sum(state * k[None, :], axis=1) + delta = (v - memory) * tl.load(beta_ptr + head_offset).to(tl.float32) + state += delta[:, None] * k[None, :] + output = tl.sum(state * q[None, :], axis=1) + tl.store( + output_ptr + head_offset * V + value_offsets, + output, + mask=value_mask, + ) + token_idx += 1 + + tl.store(final_state_ptr + state_offsets, state, mask=state_mask) + + +def _qwen35_gdn_prefill_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + g = g.float().contiguous() + beta = beta.float().contiguous() + initial_state = initial_state.float().contiguous() + cu_seqlens = cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + output = torch.empty_like(v) + final_state = torch.empty_like(initial_state) + key_dim = q.shape[-1] + value_dim = v.shape[-1] + block_key = triton.next_power_of_2(key_dim) + block_value = min(32, triton.next_power_of_2(value_dim)) + grid = ( + initial_state.shape[0], + v.shape[1], + triton.cdiv(value_dim, block_value), + ) + _qwen35_gdn_prefill_kernel[grid]( + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + output, + final_state, + H_Q=q.shape[1], + H_V=v.shape[1], + K=key_dim, + V=value_dim, + BK=block_key, + BV=block_value, + num_warps=4, + num_stages=2, + ) + return output, final_state + + +@triton.jit +def _qwen35_gdn_decode_kernel( + q_ptr, + k_ptr, + v_ptr, + a_ptr, + b_ptr, + A_log_ptr, + dt_bias_ptr, + state_ptr, + state_indices_ptr, + positions_ptr, + output_ptr, + H_Q: tl.constexpr, + H_V: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # The value block is the fastest-varying grid axis so that consecutive CTAs + # walk one request's recurrent state contiguously. With the token axis first, + # a scheduled wave touched one 16 KiB slice from each of B slots that sit + # 4 MiB apart, and this bandwidth-bound kernel lost ~13% to the scatter. + value_block = tl.program_id(0) + token_idx = tl.program_id(1) + value_head = tl.program_id(2) + key_head = value_head // (H_V // H_Q) + + key_offsets = tl.arange(0, BK) + value_offsets = value_block * BV + tl.arange(0, BV) + key_mask = key_offsets < K + value_mask = value_offsets < V + + q = tl.load( + q_ptr + (token_idx * H_Q + key_head) * K + key_offsets, + mask=key_mask, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_ptr + (token_idx * H_Q + key_head) * K + key_offsets, + mask=key_mask, + other=0.0, + ).to(tl.float32) + q *= tl.rsqrt(tl.sum(q * q, axis=0) + 1e-6) * (K**-0.5) + k *= tl.rsqrt(tl.sum(k * k, axis=0) + 1e-6) + + slot = tl.load(state_indices_ptr + token_idx).to(tl.int64) + position = tl.load(positions_ptr + token_idx) + state_offsets = ( + ((slot * H_V + value_head) * V + value_offsets[:, None]) * K + + key_offsets[None, :] + ) + state_mask = value_mask[:, None] & key_mask[None, :] + state = tl.load(state_ptr + state_offsets, mask=state_mask, other=0.0).to( + tl.float32 + ) + state = tl.where(position == 0, 0.0, state) + + head_offset = token_idx * H_V + value_head + a = tl.load(a_ptr + head_offset).to(tl.float32) + b = tl.load(b_ptr + head_offset).to(tl.float32) + A_log = tl.load(A_log_ptr + value_head).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + value_head).to(tl.float32) + softplus_arg = a + dt_bias + softplus = tl.where( + softplus_arg <= 20.0, + tl.log(1.0 + tl.exp(softplus_arg)), + softplus_arg, + ) + state *= tl.exp(-tl.exp(A_log) * softplus) + + v = tl.load( + v_ptr + (token_idx * H_V + value_head) * V + value_offsets, + mask=value_mask, + other=0.0, + ).to(tl.float32) + memory = tl.sum(state * k[None, :], axis=1) + delta = (v - memory) * tl.sigmoid(b) + state += delta[:, None] * k[None, :] + output = tl.sum(state * q[None, :], axis=1) + + tl.store(state_ptr + state_offsets, state, mask=state_mask) + tl.store( + output_ptr + (token_idx * H_V + value_head) * V + value_offsets, + output, + mask=value_mask, + ) + + +def qwen35_gdn_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """Fused single-token GDN update; mutates selected recurrent-state slots.""" + + if not q.is_cuda: + outputs = [] + for token_idx in range(q.shape[0]): + slot = int(state_indices[token_idx]) + initial = state_pool[slot] + if int(positions[token_idx]) == 0: + initial = torch.zeros_like(initial) + g, beta = qwen35_gdn_gates( + a[token_idx], b[token_idx], A_log, dt_bias + ) + output, final_state = qwen35_recurrent_gated_delta_rule( + q[token_idx : token_idx + 1], + k[token_idx : token_idx + 1], + v[token_idx : token_idx + 1], + g.unsqueeze(0), + beta.unsqueeze(0), + initial, + ) + state_pool[slot].copy_(final_state) + outputs.append(output) + return torch.cat(outputs, dim=0) + + # q/k/v arrive as strided views of the fused conv output. Reading them in place + # by stride was measured and REJECTED: it removes 108 copy launches/step worth + # ~0.62 ms but costs the bandwidth-bound recurrent kernel 25.65 us per launch + # (+0.92 ms/step), because materializing them also warms L2 for the read that + # immediately follows. Net +0.30 ms/step, so the copies stay. + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + a = a.contiguous() + b = b.contiguous() + output = torch.empty_like(v) + key_dim = q.shape[-1] + value_dim = v.shape[-1] + block_key = triton.next_power_of_2(key_dim) + block_value = min(32, triton.next_power_of_2(value_dim)) + grid = (triton.cdiv(value_dim, block_value), q.shape[0], v.shape[1]) + _qwen35_gdn_decode_kernel[grid]( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state_pool, + state_indices, + positions, + output, + H_Q=q.shape[1], + H_V=v.shape[1], + K=key_dim, + V=value_dim, + BK=block_key, + BV=block_value, + # Measured across B=32..512: 8 warps beat 4 by 7-12%. This is a + # read-modify-write of a 4 MiB/request state, so it wants issue width to + # cover memory latency rather than a narrow tile. + num_warps=8, + num_stages=2, + ) + return output + + +__all__ = [ + "qwen35_causal_conv1d_decode", + "qwen35_causal_conv1d_prefill", + "qwen35_gdn_decode", + "qwen35_gdn_gates", + "qwen35_gdn_prefill", + "qwen35_l2norm", + "qwen35_recurrent_gated_delta_rule", + "qwen35_recurrent_gated_delta_rule_packed", + "qwen35_rmsnorm_gated", +] diff --git a/python/minisgl/kvcache/__init__.py b/python/minisgl/kvcache/__init__.py index 3f5390c..e0bed93 100644 --- a/python/minisgl/kvcache/__init__.py +++ b/python/minisgl/kvcache/__init__.py @@ -41,9 +41,17 @@ def create_kvcache_pool( head_dim=model_config.head_dim, device=device, dtype=dtype, + layer_ids=tuple( + idx + for idx, layer_type in enumerate(model_config.layer_types) + if layer_type == "full_attention" + ), ) +from .hybrid_pool import HybridGDNStatePool, create_hybrid_gdn_state_pool + + @SUPPORTED_CACHE_MANAGER.register("naive") def create_naive_cache(device: torch.device): from .naive_cache import NaivePrefixCache @@ -71,4 +79,6 @@ def create_prefix_cache(device: torch.device, type: str) -> BasePrefixCache: "SizeInfo", "MatchResult", "SUPPORTED_CACHE_MANAGER", + "HybridGDNStatePool", + "create_hybrid_gdn_state_pool", ] diff --git a/python/minisgl/kvcache/hybrid_pool.py b/python/minisgl/kvcache/hybrid_pool.py new file mode 100644 index 0000000..537d23d --- /dev/null +++ b/python/minisgl/kvcache/hybrid_pool.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from minisgl.distributed import get_tp_info +from minisgl.utils import div_even + +if TYPE_CHECKING: + from minisgl.models import ModelConfig + + +@dataclass(frozen=True) +class GDNLayerState: + conv: torch.Tensor + recurrent: torch.Tensor + + +class HybridGDNStatePool: + """Fixed request-slot state for all Qwen3.5 linear-attention layers.""" + + def __init__( + self, + *, + layer_ids: tuple[int, ...], + max_slots: int, + conv_channels: int, + conv_width: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + conv_dtype: torch.dtype, + state_dtype: torch.dtype, + device: torch.device, + ) -> None: + self.layer_ids = layer_ids + self._layer_to_index = {layer_id: idx for idx, layer_id in enumerate(layer_ids)} + self.max_slots = int(max_slots) + self.conv_state = torch.zeros( + len(layer_ids), + max_slots, + conv_channels, + conv_width - 1, + dtype=conv_dtype, + device=device, + ) + self.recurrent_state = torch.zeros( + len(layer_ids), + max_slots, + num_value_heads, + value_head_dim, + key_head_dim, + dtype=state_dtype, + device=device, + ) + + def layer_state(self, layer_id: int) -> GDNLayerState: + try: + index = self._layer_to_index[int(layer_id)] + except KeyError as exc: + raise KeyError(f"layer {layer_id} is not a GDN layer") from exc + return GDNLayerState(self.conv_state[index], self.recurrent_state[index]) + + def clear_slots(self, slots: torch.Tensor) -> None: + self.conv_state.index_fill_(1, slots, 0) + self.recurrent_state.index_fill_(1, slots, 0) + + @property + def bytes_per_slot(self) -> int: + return ( + self.conv_state[:, 0].numel() * self.conv_state.element_size() + + self.recurrent_state[:, 0].numel() + * self.recurrent_state.element_size() + ) + + +def create_hybrid_gdn_state_pool( + model_config: ModelConfig, + *, + max_slots: int, + dtype: torch.dtype, + device: torch.device, +) -> HybridGDNStatePool | None: + layer_ids = tuple( + idx + for idx, layer_type in enumerate(model_config.layer_types) + if layer_type == "linear_attention" + ) + if not layer_ids: + return None + tp_size = get_tp_info().size + local_key_heads = div_even(model_config.linear_num_key_heads, tp_size) + local_value_heads = div_even(model_config.linear_num_value_heads, tp_size) + conv_channels = ( + 2 * local_key_heads * model_config.linear_key_head_dim + + local_value_heads * model_config.linear_value_head_dim + ) + state_dtype = ( + torch.float32 + if model_config.mamba_ssm_dtype == "float32" + else dtype + ) + return HybridGDNStatePool( + layer_ids=layer_ids, + max_slots=max_slots, + conv_channels=conv_channels, + conv_width=model_config.linear_conv_kernel_dim, + num_value_heads=local_value_heads, + key_head_dim=model_config.linear_key_head_dim, + value_head_dim=model_config.linear_value_head_dim, + conv_dtype=dtype, + state_dtype=state_dtype, + device=device, + ) + + +__all__ = ["GDNLayerState", "HybridGDNStatePool", "create_hybrid_gdn_state_pool"] diff --git a/python/minisgl/kvcache/mha_pool.py b/python/minisgl/kvcache/mha_pool.py index f9a681f..9308dff 100644 --- a/python/minisgl/kvcache/mha_pool.py +++ b/python/minisgl/kvcache/mha_pool.py @@ -23,11 +23,16 @@ def __init__( page_size: int, dtype: torch.dtype, device: torch.device, + layer_ids: tuple[int, ...] | None = None, ) -> None: tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) + self.layer_ids = tuple(range(num_layers)) if layer_ids is None else layer_ids + self._layer_to_storage = { + layer_id: storage_id for storage_id, layer_id in enumerate(self.layer_ids) + } self._kv_buffer = torch.empty( - (2, num_layers, num_pages, page_size, local_kv_heads, head_dim), + (2, len(self.layer_ids), num_pages, page_size, local_kv_heads, head_dim), device=device, dtype=dtype, ) @@ -37,10 +42,10 @@ def __init__( self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) def k_cache(self, index: int) -> torch.Tensor: - return self._k_buffer[index] + return self._k_buffer[self._storage_index(index)] def v_cache(self, index: int) -> torch.Tensor: - return self._v_buffer[index] + return self._v_buffer[self._storage_index(index)] def store_kv( self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int @@ -48,8 +53,8 @@ def store_kv( from minisgl.kernel import store_cache store_cache( - k_cache=self._k_buffer[layer_id].view(self._storage_shape), - v_cache=self._v_buffer[layer_id].view(self._storage_shape), + k_cache=self._k_buffer[self._storage_index(layer_id)].view(self._storage_shape), + v_cache=self._v_buffer[self._storage_index(layer_id)].view(self._storage_shape), indices=out_loc, k=k, v=v, @@ -62,3 +67,9 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: return self._kv_buffer.dtype + + def _storage_index(self, layer_id: int) -> int: + try: + return self._layer_to_storage[int(layer_id)] + except KeyError as exc: + raise KeyError(f"layer {layer_id} has no full-attention KV cache") from exc diff --git a/python/minisgl/layers/__init__.py b/python/minisgl/layers/__init__.py index 8a34da0..4627e3d 100644 --- a/python/minisgl/layers/__init__.py +++ b/python/minisgl/layers/__init__.py @@ -5,12 +5,13 @@ from .linear import ( LinearColParallelMerged, LinearOProj, + LinearQKVGatedMerged, LinearQKVMerged, LinearReplicated, LinearRowParallel, ) from .moe import MoELayer -from .norm import RMSNorm, RMSNormCrossHead, RMSNormFused +from .norm import RMSNorm, RMSNormCrossHead, RMSNormFused, RMSNormOffset, RMSNormOffsetFused from .rotary import get_rope, set_rope_device __all__ = [ @@ -27,9 +28,12 @@ "LinearRowParallel", "LinearOProj", "LinearQKVMerged", + "LinearQKVGatedMerged", "RMSNorm", "RMSNormFused", "RMSNormCrossHead", + "RMSNormOffset", + "RMSNormOffsetFused", "get_rope", "set_rope_device", "LinearReplicated", diff --git a/python/minisgl/layers/attention.py b/python/minisgl/layers/attention.py index 7f57599..e1f081a 100644 --- a/python/minisgl/layers/attention.py +++ b/python/minisgl/layers/attention.py @@ -9,7 +9,7 @@ from minisgl.utils import div_even from .base import StateLessOP -from .rotary import get_rope +from .rotary import freeze_rope_scaling, get_rope if TYPE_CHECKING: from minisgl.layers import RMSNorm @@ -40,7 +40,7 @@ def __init__( rotary_dim=rotary_config.rotary_dim, max_position=rotary_config.max_position, base=rotary_config.base, - rope_scaling=tuple(rotary_config.scaling.items()) if rotary_config.scaling else None, + rope_scaling=freeze_rope_scaling(rotary_config.scaling), ) self.q_norm = q_norm self.k_norm = k_norm diff --git a/python/minisgl/layers/linear.py b/python/minisgl/layers/linear.py index 7da63b8..4090727 100644 --- a/python/minisgl/layers/linear.py +++ b/python/minisgl/layers/linear.py @@ -334,6 +334,32 @@ def __init__( self._use_bf16_cublas_for_fp8 = False +class LinearQKVGatedMerged(_LinearTPImpl): + """QKV projection whose Q rows carry per-head query and output gate.""" + + def __init__( + self, + hidden_size: int, + head_dim: int, + num_qo_heads: int, + num_kv_heads: int, + has_bias: bool, + quant: Any | None = None, + ): + tp_info = get_tp_info() + local_num_qo = div_even(num_qo_heads, tp_info.size) + local_num_kv = div_even(num_kv_heads, tp_info.size, allow_replicate=True) + super().__init__( + hidden_size, + (2 * num_qo_heads + 2 * num_kv_heads) * head_dim, + hidden_size, + (2 * local_num_qo + 2 * local_num_kv) * head_dim, + has_bias, + quant=quant, + ) + self._use_bf16_cublas_for_fp8 = False + + class LinearOProj(_LinearTPImpl): def __init__( self, diff --git a/python/minisgl/layers/moe/layer.py b/python/minisgl/layers/moe/layer.py index ffa1a9f..0ee8c93 100644 --- a/python/minisgl/layers/moe/layer.py +++ b/python/minisgl/layers/moe/layer.py @@ -72,6 +72,7 @@ def __init__( if ep_size == 1: self.local_num_experts = num_experts + self.local_expert_offset = 0 self._expert_map_dev: torch.Tensor | None = None intermediate_size_per_partition = div_even(intermediate_size, tp_info.size) else: @@ -79,6 +80,7 @@ def __init__( self.local_num_experts, expert_map_cpu = _build_expert_map( num_experts, ep_info.rank, ep_size ) + self.local_expert_offset = int(ep_info.rank) * self.local_num_experts device = ( torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() @@ -90,8 +92,48 @@ def __init__( intermediate_size_per_partition = intermediate_size self._is_fp8 = quant is not None and quant.method in ("fp8", "fp8_channel") + self._is_nvfp4 = quant is not None and quant.method == "modelopt_nvfp4" self._fp8_scale_format = None - if self._is_fp8: + if self._is_nvfp4: + if hidden_size % 16 or intermediate_size_per_partition % 16: + raise RuntimeError( + "ModelOpt NVFP4 requires hidden and local intermediate sizes " + f"divisible by 16, got {hidden_size} and " + f"{intermediate_size_per_partition}" + ) + self.gate_up_proj = torch.empty( + self.local_num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 2, + dtype=torch.uint8, + ) + self.gate_up_proj_scale = torch.empty( + self.local_num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 16, + dtype=torch.float8_e4m3fn, + ) + self.gate_up_proj_weight_scale_2 = torch.empty( + self.local_num_experts, 2, dtype=torch.float32 + ) + self.gate_up_proj_input_scale = torch.empty(2, dtype=torch.float32) + self.down_proj = torch.empty( + self.local_num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=torch.uint8, + ) + self.down_proj_scale = torch.empty( + self.local_num_experts, + hidden_size, + intermediate_size_per_partition // 16, + dtype=torch.float8_e4m3fn, + ) + self.down_proj_weight_scale_2 = torch.empty( + self.local_num_experts, dtype=torch.float32 + ) + self.down_proj_input_scale = torch.empty((), dtype=torch.float32) + elif self._is_fp8: assert quant is not None if quant.method == "fp8_channel": block_out, block_in = (1, 1) @@ -168,6 +210,14 @@ def __init__( if a2a_backend is not None else getattr(server_args, "afd_moe_a2a_backend", "none") ) + if self._is_nvfp4: + if runner_backend == "auto": + runner_backend = "flashinfer_nvfp4" + elif runner_backend != "flashinfer_nvfp4": + raise RuntimeError( + "ModelOpt NVFP4 routed experts require " + "afd_moe_runner_backend=flashinfer_nvfp4" + ) if a2a_backend == "deepep" and runner_backend == "auto": runner_backend = "deep_gemm" runner_config = MoeRunnerConfig( @@ -192,8 +242,13 @@ def __init__( moe_runner_config=runner_config, dp_size=self.moe_dp_size, ) - if isinstance(self.dispatcher, DeepEPDispatcher) and runner_backend != "deep_gemm": - raise RuntimeError("DeepEP elastic MoE requires afd_moe_runner_backend=deep_gemm") + if isinstance(self.dispatcher, DeepEPDispatcher) and runner_backend not in { + "deep_gemm", + "flashinfer_nvfp4", + }: + raise RuntimeError( + "DeepEP elastic MoE requires deep_gemm or flashinfer_nvfp4" + ) runner_config.runner_backend = runner_backend self.runner = create_moe_runner( runner_backend=runner_backend, diff --git a/python/minisgl/layers/moe/moe_runner/__init__.py b/python/minisgl/layers/moe/moe_runner/__init__.py index 492d3ab..9768073 100644 --- a/python/minisgl/layers/moe/moe_runner/__init__.py +++ b/python/minisgl/layers/moe/moe_runner/__init__.py @@ -2,6 +2,7 @@ from .base import MoeA2ABackend, MoeRunner, MoeRunnerBackend, MoeRunnerConfig from .deepgemm_grouped import DeepGEMMGroupedRunner +from .flashinfer_nvfp4 import FlashInferNvFp4Runner from .triton import TritonRunner from .triton_fp8 import TritonFp8Runner @@ -17,6 +18,8 @@ def create_moe_runner( return TritonFp8Runner(config=moe_runner_config) if backend.is_deep_gemm(): return DeepGEMMGroupedRunner(config=moe_runner_config) + if backend.is_flashinfer_nvfp4(): + return FlashInferNvFp4Runner(config=moe_runner_config) if backend.is_triton_kernels(): raise NotImplementedError( "runner_backend='triton_kernel' is not implemented in the no-DeepEP EP port" @@ -30,6 +33,7 @@ def create_moe_runner( "MoeRunnerBackend", "MoeRunnerConfig", "DeepGEMMGroupedRunner", + "FlashInferNvFp4Runner", "TritonRunner", "TritonFp8Runner", "create_moe_runner", diff --git a/python/minisgl/layers/moe/moe_runner/base.py b/python/minisgl/layers/moe/moe_runner/base.py index 919ccd7..fc6e93d 100644 --- a/python/minisgl/layers/moe/moe_runner/base.py +++ b/python/minisgl/layers/moe/moe_runner/base.py @@ -34,6 +34,7 @@ class MoeRunnerBackend(Enum): TRITON_FP8 = "triton_fp8" TRITON_KERNELS = "triton_kernel" DEEP_GEMM = "deep_gemm" + FLASHINFER_NVFP4 = "flashinfer_nvfp4" def is_auto(self) -> bool: return self == MoeRunnerBackend.AUTO @@ -50,6 +51,9 @@ def is_triton_kernels(self) -> bool: def is_deep_gemm(self) -> bool: return self == MoeRunnerBackend.DEEP_GEMM + def is_flashinfer_nvfp4(self) -> bool: + return self == MoeRunnerBackend.FLASHINFER_NVFP4 + @dataclass class MoeRunnerConfig: diff --git a/python/minisgl/layers/moe/moe_runner/flashinfer_nvfp4.py b/python/minisgl/layers/moe/moe_runner/flashinfer_nvfp4.py new file mode 100644 index 0000000..e867472 --- /dev/null +++ b/python/minisgl/layers/moe/moe_runner/flashinfer_nvfp4.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from minisgl.kernel.nvfp4_moe import ( + NvFp4PreparedMoEWeights, + prepare_nvfp4_moe_weights, + run_nvfp4_moe, +) + +from .base import MoeRunner, MoeRunnerBackend, MoeRunnerConfig + + +def _deepep_valid_rows(dispatch_output: Any) -> torch.Tensor | None: + handle = getattr(dispatch_output, "handle", None) + runtime_handle = getattr(handle, "handle", None) + psum = getattr(runtime_handle, "psum_num_recv_tokens_per_scaleup_rank", None) + if psum is None: + return None + rows = torch.arange( + dispatch_output.hidden_states.shape[0], + device=dispatch_output.hidden_states.device, + dtype=psum.dtype, + ) + return rows < psum[-1] + + +def normalize_deepep_nvfp4_topk( + dispatch_output: Any, + layer: Any, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert DeepEP rank-local routes into the model-global NVFP4 contract. + + Non-expanded DeepEP output is graph-static: rows beyond the device-side + receive count are allocation padding. Mask those rows without a host + ``item()`` synchronization, preserve ``-1`` route sentinels, and translate + valid receiving-rank-local expert ids to the global ids expected by the + FlashInfer EP launcher. + """ + + topk_ids = getattr(dispatch_output, "topk_ids", None) + topk_weights = getattr(dispatch_output, "topk_weights", None) + if topk_ids is None or topk_weights is None: + raise RuntimeError("NVFP4 DeepEP dispatch requires explicit top-k metadata") + if topk_ids.ndim != 2 or topk_weights.shape != topk_ids.shape: + raise RuntimeError( + "NVFP4 DeepEP top-k shape mismatch: " + f"ids={tuple(topk_ids.shape)} weights={tuple(topk_weights.shape)}" + ) + + ids = topk_ids.to(dtype=torch.int32) + valid = ids >= 0 + local_num_experts = int(layer.local_num_experts) + valid &= ids < local_num_experts + valid_rows = _deepep_valid_rows(dispatch_output) + if valid_rows is not None: + valid &= valid_rows[:, None] + global_ids = torch.where( + valid, + ids + int(layer.local_expert_offset), + torch.full_like(ids, -1), + ).contiguous() + weights = torch.where( + valid, + topk_weights.to(dtype=torch.float32), + torch.zeros((), device=topk_weights.device, dtype=torch.float32), + ).contiguous() + return weights, global_ids + + +class FlashInferNvFp4Runner(MoeRunner): + """Blackwell native ModelOpt NVFP4 routed-expert runner.""" + + @property + def runner_backend(self) -> MoeRunnerBackend: + return MoeRunnerBackend.FLASHINFER_NVFP4 + + @staticmethod + def _topk(dispatch_output: Any, layer: Any) -> tuple[torch.Tensor, torch.Tensor]: + topk_output = getattr(dispatch_output, "topk_output", None) + if isinstance(topk_output, tuple) and len(topk_output) == 2: + return topk_output + topk_ids = getattr(dispatch_output, "topk_ids", None) + topk_weights = getattr(dispatch_output, "topk_weights", None) + if topk_ids is not None and topk_weights is not None: + if _deepep_valid_rows(dispatch_output) is not None: + return normalize_deepep_nvfp4_topk(dispatch_output, layer) + return topk_weights, topk_ids + if isinstance(topk_output, torch.Tensor): + from minisgl.moe.fused import fused_topk + + return fused_topk( + hidden_states=dispatch_output.hidden_states, + gating_output=topk_output, + topk=layer.top_k, + renormalize=layer.renormalize, + ) + raise RuntimeError("NVFP4 runner requires router logits or explicit top-k tensors") + + @staticmethod + def _prepare(layer: Any) -> NvFp4PreparedMoEWeights: + cached = getattr(layer, "_nvfp4_prepared", None) + if cached is not None: + return cached + prepared = prepare_nvfp4_moe_weights( + layer.gate_up_proj, + layer.gate_up_proj_scale, + layer.gate_up_proj_weight_scale_2, + layer.gate_up_proj_input_scale, + layer.down_proj, + layer.down_proj_scale, + layer.down_proj_weight_scale_2, + layer.down_proj_input_scale, + global_num_experts=layer.num_experts, + local_expert_offset=layer.local_expert_offset, + ) + # Replace serialized tensors with their native reordered equivalents, + # then release all now-fused global-scale inputs. The private dataclass + # owns the remaining FP32 scale vectors used by the kernel. + layer.gate_up_proj = prepared.gate_up_proj + layer.gate_up_proj_scale = prepared.gate_up_proj_scale + layer.down_proj = prepared.down_proj + layer.down_proj_scale = prepared.down_proj_scale + layer.gate_up_proj_weight_scale_2 = None + layer.gate_up_proj_input_scale = None + layer.down_proj_weight_scale_2 = None + layer.down_proj_input_scale = None + layer._nvfp4_prepared = prepared + return prepared + + def prewarm(self, layer: Any, *, dtype: torch.dtype = torch.bfloat16) -> None: + """Prepare one layer and compile/launch the shared MoE kernel before serving.""" + + prepared = self._prepare(layer) + top_k = int(self.config.top_k or layer.top_k) + local_num_experts = int(layer.local_num_experts) + if top_k > local_num_experts: + raise RuntimeError( + "NVFP4 prewarm requires at least top_k local experts: " + f"top_k={top_k} local_num_experts={local_num_experts}" + ) + device = prepared.gate_up_proj.device + hidden_states = torch.zeros( + (1, prepared.hidden_size), device=device, dtype=dtype + ) + topk_ids = ( + torch.arange(top_k, device=device, dtype=torch.int32) + + int(layer.local_expert_offset) + ).unsqueeze(0) + topk_weights = torch.full( + (1, top_k), 1.0 / top_k, device=device, dtype=torch.float32 + ) + output = run_nvfp4_moe( + prepared, + hidden_states, + topk_ids, + topk_weights, + ) + if device.type == "cuda": + torch.cuda.synchronize(device) + del output + + def apply(self, dispatch_output: Any, layer: Any) -> torch.Tensor: + hidden_states = dispatch_output.hidden_states + if hidden_states.shape[0] == 0: + return hidden_states + valid_rows = _deepep_valid_rows(dispatch_output) + if valid_rows is not None: + hidden_states = torch.where( + valid_rows[:, None], + hidden_states, + torch.zeros((), device=hidden_states.device, dtype=hidden_states.dtype), + ) + topk_weights, topk_ids = self._topk(dispatch_output, layer) + prepared = self._prepare(layer) + return run_nvfp4_moe( + prepared, + hidden_states, + topk_ids, + topk_weights, + ) + + +__all__ = ["FlashInferNvFp4Runner", "normalize_deepep_nvfp4_topk"] diff --git a/python/minisgl/layers/moe/token_dispatcher/deepep.py b/python/minisgl/layers/moe/token_dispatcher/deepep.py index 6fddd88..f04f4db 100644 --- a/python/minisgl/layers/moe/token_dispatcher/deepep.py +++ b/python/minisgl/layers/moe/token_dispatcher/deepep.py @@ -62,6 +62,9 @@ def __init__(self, moe_runner_config: MoeRunnerConfig): self.num_experts = int(moe_runner_config.num_experts or 0) self.num_local_experts = int(moe_runner_config.num_local_experts or 0) self.top_k = int(moe_runner_config.top_k or 0) + self.use_expanded_layout = ( + str(moe_runner_config.runner_backend) != "flashinfer_nvfp4" + ) self._group_expert_map: torch.Tensor | None = None @staticmethod @@ -159,8 +162,13 @@ def dispatch( group_topk_ids.contiguous(), topk_weights.contiguous(), hidden_states_scale=dispatch_x_scale, - expert_alignment=self._deepgemm_expert_alignment(), + expert_alignment=( + self._deepgemm_expert_alignment() + if self.use_expanded_layout + else 1 + ), num_max_dispatch_tokens_per_rank=ctx.moe_deepep_dispatch_max_tokens_per_rank, + use_expanded_layout=self.use_expanded_layout, ) return DeepEPElasticDispatchOutput( hidden_states=recv_x, @@ -214,6 +222,34 @@ def _post_permute_runner_to_deepep_elastic( ) +@register_pre_permute("deepep_elastic", "flashinfer_nvfp4") +def _pre_permute_deepep_elastic_to_flashinfer_nvfp4( + dispatch_output: DeepEPElasticDispatchOutput, +) -> DeepEPElasticDispatchOutput: + if dispatch_output.topk_ids is None or dispatch_output.topk_weights is None: + raise RuntimeError( + "FlashInfer NVFP4 requires DeepEP non-expanded top-k metadata" + ) + if dispatch_output.topk_ids.ndim != 2 or dispatch_output.topk_weights.ndim != 2: + raise RuntimeError( + "FlashInfer NVFP4 requires 2D DeepEP top-k ids and weights" + ) + if dispatch_output.hidden_states_scale is not None: + raise RuntimeError("FlashInfer NVFP4 DeepEP transport requires BF16 activations") + return dispatch_output + + +@register_post_permute("flashinfer_nvfp4", "deepep_elastic") +def _post_permute_flashinfer_nvfp4_to_deepep_elastic( + runner_output: torch.Tensor, + dispatch_output: DeepEPElasticDispatchOutput, +) -> DeepEPElasticCombineInput: + return DeepEPElasticCombineInput( + hidden_states=runner_output, + handle=dispatch_output.handle, + ) + + __all__ = [ "DeepEPDispatcher", "DeepEPElasticCombineInput", diff --git a/python/minisgl/layers/moe/token_dispatcher/standard.py b/python/minisgl/layers/moe/token_dispatcher/standard.py index 8d76b5f..7c9570c 100644 --- a/python/minisgl/layers/moe/token_dispatcher/standard.py +++ b/python/minisgl/layers/moe/token_dispatcher/standard.py @@ -122,4 +122,20 @@ def _post_permute_triton_fp8_to_standard( return StandardCombineInput(hidden_states=runner_output) +@register_pre_permute("standard", "flashinfer_nvfp4") +def _pre_permute_standard_to_flashinfer_nvfp4( + dispatch_output: StandardDispatchOutput, +) -> StandardDispatchOutput: + return dispatch_output + + +@register_post_permute("flashinfer_nvfp4", "standard") +def _post_permute_flashinfer_nvfp4_to_standard( + runner_output: torch.Tensor, + dispatch_output: StandardDispatchOutput, +) -> StandardCombineInput: + del dispatch_output + return StandardCombineInput(hidden_states=runner_output) + + __all__ = ["StandardDispatchOutput", "StandardCombineInput", "StandardDispatcher"] diff --git a/python/minisgl/layers/norm.py b/python/minisgl/layers/norm.py index d8d82a8..69480c1 100644 --- a/python/minisgl/layers/norm.py +++ b/python/minisgl/layers/norm.py @@ -38,6 +38,34 @@ def forward( return x, residual +class RMSNormOffset(RMSNorm): + """Gemma/Qwen3.5 RMSNorm with checkpoint multiplier ``1 + weight``. + + The offset is folded once while loading so existing FlashInfer and fused + QK-norm kernels consume the effective multiplier without a per-token add. + """ + + def load_state_dict(self, state_dict, *, prefix="", _internal=False) -> None: + key = f"{prefix}.weight" if prefix else "weight" + item = state_dict.pop(key) + assert self.weight.shape == item.shape and self.weight.dtype == item.dtype + self.weight = item + 1 + if not _internal and state_dict: + raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}") + + +class RMSNormOffsetFused(RMSNormFused): + """Residual-fused form of :class:`RMSNormOffset`.""" + + def load_state_dict(self, state_dict, *, prefix="", _internal=False) -> None: + key = f"{prefix}.weight" if prefix else "weight" + item = state_dict.pop(key) + assert self.weight.shape == item.shape and self.weight.dtype == item.dtype + self.weight = item + 1 + if not _internal and state_dict: + raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}") + + def _qk_variance(q: torch.Tensor, k: torch.Tensor): q = q.float() k = k.float() diff --git a/python/minisgl/layers/rotary.py b/python/minisgl/layers/rotary.py index f715dbd..9ee1c64 100644 --- a/python/minisgl/layers/rotary.py +++ b/python/minisgl/layers/rotary.py @@ -126,6 +126,23 @@ def set_rope_device(device: torch.device): _ROPE_DEVICE = device +def freeze_rope_scaling( + rope_scaling: Dict[str, Any] | None, +) -> Tuple[Tuple[str, Any], ...] | None: + """Build a recursively hashable RoPE cache key from HF config values.""" + + def freeze(value: Any) -> Any: + if isinstance(value, dict): + return tuple((key, freeze(item)) for key, item in value.items()) + if isinstance(value, (list, tuple)): + return tuple(freeze(item) for item in value) + return value + + if rope_scaling is None: + return None + return tuple((key, freeze(value)) for key, value in rope_scaling.items()) + + @functools.cache def get_rope( head_dim: int, @@ -147,4 +164,9 @@ def get_rope( return _get_rope(head_dim, rotary_dim, max_position, base, rope_map) -__all__ = ["get_rope", "RotaryEmbedding", "set_rope_device"] +__all__ = [ + "freeze_rope_scaling", + "get_rope", + "RotaryEmbedding", + "set_rope_device", +] diff --git a/python/minisgl/models/afd.py b/python/minisgl/models/afd.py index e146dec..f29836b 100644 --- a/python/minisgl/models/afd.py +++ b/python/minisgl/models/afd.py @@ -25,10 +25,34 @@ def extract_stage_state_dict( stage_model: BaseOP, full_state_dict: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]: - expected_keys = stage_model.state_dict().keys() + expected_state = stage_model.state_dict() + expected_keys = expected_state.keys() missing = [key for key in expected_keys if key not in full_state_dict] if missing: raise RuntimeError(f"Missing keys for AFD model stage: {missing[:8]}") + incompatible = [ + ( + key, + tuple(expected_state[key].shape), + expected_state[key].dtype, + tuple(full_state_dict[key].shape), + full_state_dict[key].dtype, + ) + for key in expected_keys + if ( + expected_state[key].shape != full_state_dict[key].shape + or expected_state[key].dtype != full_state_dict[key].dtype + ) + ] + if incompatible: + details = [ + ( + f"{key}: expected shape={expected_shape} dtype={expected_dtype}, " + f"got shape={actual_shape} dtype={actual_dtype}" + ) + for key, expected_shape, expected_dtype, actual_shape, actual_dtype in incompatible[:8] + ] + raise RuntimeError("Incompatible AFD stage tensors: " + " | ".join(details)) return {key: full_state_dict[key] for key in expected_keys} diff --git a/python/minisgl/models/config.py b/python/minisgl/models/config.py index d6f015a..7b04909 100644 --- a/python/minisgl/models/config.py +++ b/python/minisgl/models/config.py @@ -1,5 +1,6 @@ from __future__ import annotations from dataclasses import dataclass, field +from fnmatch import fnmatchcase from typing import Any, Dict, Optional, Tuple from transformers import PretrainedConfig @@ -20,6 +21,22 @@ class QuantConfig: method: str activation_scheme: str weight_block_size: Tuple[int, int] + group_size: int | None = None + ignored_modules: Tuple[str, ...] = () + kv_cache_dtype: str | None = None + + def is_ignored(self, module_name: str) -> bool: + candidates = [module_name] + language_prefix = "model.language_model." + if module_name.startswith(language_prefix): + candidates.append("model." + module_name.removeprefix(language_prefix)) + elif module_name.startswith("model."): + candidates.append(language_prefix + module_name.removeprefix("model.")) + return any( + fnmatchcase(candidate, pattern) + for pattern in self.ignored_modules + for candidate in candidates + ) @dataclass(frozen=True) @@ -49,42 +66,113 @@ class ModelConfig: rope_scaling: Dict[str, Any] | None = None rope_theta: float = 10000.0 partial_rotary_factor: float | None = None + layer_types: Tuple[str, ...] = () + attn_output_gate: bool = False + linear_conv_kernel_dim: int = 0 + linear_key_head_dim: int = 0 + linear_value_head_dim: int = 0 + linear_num_key_heads: int = 0 + linear_num_value_heads: int = 0 + shared_expert_intermediate_size: int = 0 + mamba_ssm_dtype: str = "float32" model_extra: Dict[str, Any] = field(default_factory=dict) @property def is_moe(self) -> bool: return self.num_experts > 0 + @property + def num_full_attention_layers(self) -> int: + return sum(layer_type == "full_attention" for layer_type in self.layer_types) + + @property + def num_linear_attention_layers(self) -> int: + return sum(layer_type == "linear_attention" for layer_type in self.layer_types) + + @property + def is_hybrid_attention(self) -> bool: + return self.num_linear_attention_layers > 0 + + def quant_for_modules(self, *module_names: str) -> QuantConfig | None: + if self.quant is None: + return None + if any(self.quant.is_ignored(name) for name in module_names): + return None + return self.quant + @classmethod def from_hf(cls, config: PretrainedConfig) -> ModelConfig: - if hasattr(config, "text_config") and config.text_config is not None: - top = config - config = config.text_config - for attr in ("architectures", "rope_theta", "rope_scaling", "rope_parameters"): - if not getattr(config, attr, None) and getattr(top, attr, None): - setattr(config, attr, getattr(top, attr)) - - num_kv_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) - head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - tie_word_embeddings = getattr(config, "tie_word_embeddings", False) - model_type = getattr(config, "model_type", "llama") - n_routed_experts = getattr(config, "n_routed_experts", None) + def value(obj: Any, name: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(name, default) + return getattr(obj, name, default) + + top = config + text_config = value(config, "text_config", None) + if text_config is not None: + config = text_config + + num_attention_heads = int(value(config, "num_attention_heads")) + hidden_size = int(value(config, "hidden_size")) + num_kv_heads = value(config, "num_key_value_heads", num_attention_heads) + head_dim = value(config, "head_dim", None) or hidden_size // num_attention_heads + tie_word_embeddings = value( + config, + "tie_word_embeddings", + value(top, "tie_word_embeddings", False), + ) + model_type = value(config, "model_type", value(top, "model_type", "llama")) + n_routed_experts = value(config, "n_routed_experts", None) if n_routed_experts is None: - num_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 0)) + num_experts = value( + config, "num_local_experts", value(config, "num_experts", 0) + ) n_routed_experts = num_experts else: num_experts = n_routed_experts - num_experts_per_tok = getattr(config, "num_experts_per_tok", 0) - moe_intermediate_size = getattr(config, "moe_intermediate_size", 0) or config.intermediate_size - norm_topk_prob = getattr(config, "norm_topk_prob", False) - use_routing_bias = getattr(config, "use_routing_bias", False) - scoring_func = getattr(config, "scoring_func", "softmax") - if model_type == "minimax_m2" and scoring_func == "sigmoid" and not hasattr(config, "norm_topk_prob"): + num_experts_per_tok = value(config, "num_experts_per_tok", 0) + moe_intermediate_size = int( + value(config, "moe_intermediate_size", 0) + or value(config, "intermediate_size", 0) + ) + intermediate_size = int( + value(config, "intermediate_size", 0) or moe_intermediate_size + ) + qwen35_model_types = {"qwen3_5_moe", "qwen3_5_moe_text"} + norm_topk_prob = value(config, "norm_topk_prob", model_type in qwen35_model_types) + use_routing_bias = value(config, "use_routing_bias", False) + scoring_func = value(config, "scoring_func", "softmax") + if ( + model_type == "minimax_m2" + and scoring_func == "sigmoid" + and value(config, "norm_topk_prob", None) is None + ): # vLLM hardcodes MiniMax M2 sigmoid-routing weights to be # renormalized; the HF config currently omits norm_topk_prob. norm_topk_prob = True - architectures = getattr(config, "architectures", ["LlamaForCausalLM"]) - use_qk_norm = getattr(config, "use_qk_norm", False) + architectures = value( + config, + "architectures", + value(top, "architectures", ["LlamaForCausalLM"]), + ) + use_qk_norm = value(config, "use_qk_norm", model_type in qwen35_model_types) + num_layers = int(value(config, "num_hidden_layers")) + layer_types_raw = value(config, "layer_types", None) + if layer_types_raw is None: + layer_types = ("full_attention",) * num_layers + else: + layer_types = tuple(str(item) for item in layer_types_raw) + if len(layer_types) != num_layers: + raise ValueError( + "layer_types length must equal num_hidden_layers: " + f"{len(layer_types)} != {num_layers}" + ) + invalid_layer_types = set(layer_types) - { + "full_attention", + "linear_attention", + } + if invalid_layer_types: + raise ValueError(f"Unsupported layer_types: {sorted(invalid_layer_types)}") model_extra: Dict[str, Any] = {} if model_type == "glm4_moe" or "Glm4MoeForCausalLM" in architectures: for attr, default in ( @@ -100,23 +188,27 @@ def from_hf(cls, config: PretrainedConfig) -> ModelConfig: # Llama/Qwen: rope_theta is a direct attr; Mistral: it's inside rope_scaling dict # MiniMax M2 also uses partial rotary via rotary_dim. - rope_scaling = getattr(config, "rope_scaling", None) - rope_parameters = getattr(config, "rope_parameters", None) + rope_scaling = value(config, "rope_scaling", value(top, "rope_scaling", None)) + rope_parameters = value( + config, "rope_parameters", value(top, "rope_parameters", None) + ) if rope_parameters is not None: - rope_theta = rope_parameters.get("rope_theta", None) or getattr(config, "rope_theta", None) + rope_theta = rope_parameters.get("rope_theta", None) or value( + config, "rope_theta", value(top, "rope_theta", None) + ) rope_scaling = rope_parameters else: - rope_theta = getattr(config, "rope_theta", None) + rope_theta = value(config, "rope_theta", value(top, "rope_theta", None)) if rope_theta is None and rope_scaling is not None: rope_theta = rope_scaling["rope_theta"] if rope_theta is None: rope_theta = 10000.0 - partial_rotary_factor = getattr(config, "partial_rotary_factor", None) + partial_rotary_factor = value(config, "partial_rotary_factor", None) if partial_rotary_factor is None and isinstance(rope_parameters, dict): partial_rotary_factor = rope_parameters.get("partial_rotary_factor") if partial_rotary_factor is None and isinstance(rope_scaling, dict): partial_rotary_factor = rope_scaling.get("partial_rotary_factor") - rotary_dim = getattr(config, "rotary_dim", None) + rotary_dim = value(config, "rotary_dim", None) if rotary_dim is None: rotary_dim = ( int(head_dim * float(partial_rotary_factor)) @@ -125,7 +217,11 @@ def from_hf(cls, config: PretrainedConfig) -> ModelConfig: ) quant: Optional[QuantConfig] = None - quant_config = getattr(config, "quantization_config", None) + quant_config = value( + config, + "quantization_config", + value(top, "quantization_config", None), + ) if quant_config is not None: quant_dict = ( quant_config @@ -139,6 +235,9 @@ def from_hf(cls, config: PretrainedConfig) -> ModelConfig: method="fp8", activation_scheme=quant_dict.get("activation_scheme", "dynamic"), weight_block_size=(int(block_size[0]), int(block_size[1])), + ignored_modules=tuple( + quant_dict.get("modules_to_not_convert", ()) + ), ) elif ( quant_method == "compressed-tensors" @@ -153,22 +252,40 @@ def from_hf(cls, config: PretrainedConfig) -> ModelConfig: activation_scheme="dynamic", weight_block_size=(1, 1), ) + elif quant_method == "modelopt" and quant_dict.get("quant_algo") == "NVFP4": + groups = quant_dict.get("config_groups", {}) + group0 = groups.get("group_0", {}) if isinstance(groups, dict) else {} + weights = group0.get("weights", {}) if isinstance(group0, dict) else {} + group_size = int(weights.get("group_size", 16)) + kv_scheme = quant_dict.get("kv_cache_scheme", {}) + quant = QuantConfig( + method="modelopt_nvfp4", + activation_scheme="static", + weight_block_size=(1, group_size), + group_size=group_size, + ignored_modules=tuple(quant_dict.get("ignore", ())), + kv_cache_dtype=( + f"fp{int(kv_scheme.get('num_bits', 8))}" + if isinstance(kv_scheme, dict) and kv_scheme + else None + ), + ) return cls( - num_layers=config.num_hidden_layers, - num_qo_heads=config.num_attention_heads, + num_layers=num_layers, + num_qo_heads=num_attention_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - hidden_size=config.hidden_size, - vocab_size=config.vocab_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - rms_norm_eps=config.rms_norm_eps, + hidden_size=hidden_size, + vocab_size=int(value(config, "vocab_size")), + intermediate_size=intermediate_size, + hidden_act=str(value(config, "hidden_act")), + rms_norm_eps=float(value(config, "rms_norm_eps")), tie_word_embeddings=tie_word_embeddings, rotary_config=RotaryConfig( head_dim=head_dim, rotary_dim=rotary_dim, - max_position=config.max_position_embeddings, + max_position=int(value(config, "max_position_embeddings")), base=rope_theta, scaling=rope_scaling, parameters=rope_parameters, @@ -188,5 +305,16 @@ def from_hf(cls, config: PretrainedConfig) -> ModelConfig: rope_scaling=rope_scaling, rope_theta=rope_theta, partial_rotary_factor=partial_rotary_factor, + layer_types=layer_types, + attn_output_gate=bool(value(config, "attn_output_gate", False)), + linear_conv_kernel_dim=int(value(config, "linear_conv_kernel_dim", 0)), + linear_key_head_dim=int(value(config, "linear_key_head_dim", 0)), + linear_value_head_dim=int(value(config, "linear_value_head_dim", 0)), + linear_num_key_heads=int(value(config, "linear_num_key_heads", 0)), + linear_num_value_heads=int(value(config, "linear_num_value_heads", 0)), + shared_expert_intermediate_size=int( + value(config, "shared_expert_intermediate_size", 0) + ), + mamba_ssm_dtype=str(value(config, "mamba_ssm_dtype", "float32")), model_extra=model_extra, ) diff --git a/python/minisgl/models/qwen35_moe.py b/python/minisgl/models/qwen35_moe.py new file mode 100644 index 0000000..5561a3c --- /dev/null +++ b/python/minisgl/models/qwen35_moe.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch +from minisgl.core import get_global_ctx +from minisgl.distributed import get_tp_info +from minisgl.kernel.qwen35_gdn import ( + qwen35_causal_conv1d_decode, + qwen35_causal_conv1d_prefill, + qwen35_gdn_decode, + qwen35_gdn_gates, + qwen35_gdn_prefill, + qwen35_rmsnorm_gated, +) +from minisgl.layers import ( + AttentionLayer, + BaseOP, + LinearColParallelMerged, + LinearOProj, + LinearQKVGatedMerged, + LinearReplicated, + LinearRowParallel, + MoELayer, + OPList, + ParallelLMHead, + RMSNorm, + RMSNormOffset, + RMSNormOffsetFused, + VocabParallelEmbedding, + silu_and_mul, +) +from minisgl.utils import div_even, nvtx_annotate + +from .base import BaseLLMModel + +if TYPE_CHECKING: + from .config import ModelConfig + + +class _Qwen3_5DepthwiseConv1d(BaseOP): + def __init__(self, channels: int, kernel_size: int): + self.weight = torch.empty(channels, 1, kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise RuntimeError("Qwen3.5 depthwise convolution requires request state") + + +class Qwen3_5LinearAttention(BaseOP): + """TP-sharded Qwen3.5 Gated DeltaNet with request-scoped recurrent state.""" + + def __init__(self, config: ModelConfig, layer_id: int): + tp_size = get_tp_info().size + key_dim = config.linear_num_key_heads * config.linear_key_head_dim + value_dim = config.linear_num_value_heads * config.linear_value_head_dim + local_key_heads = div_even(config.linear_num_key_heads, tp_size) + local_value_heads = div_even(config.linear_num_value_heads, tp_size) + local_conv_dim = ( + 2 * local_key_heads * config.linear_key_head_dim + + local_value_heads * config.linear_value_head_dim + ) + module_prefix = f"model.layers.{layer_id}.linear_attn" + + self.in_proj_qkv = LinearColParallelMerged( + config.hidden_size, + [key_dim, key_dim, value_dim], + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.in_proj_qkv"), + ) + self.in_proj_z = LinearColParallelMerged( + config.hidden_size, + [value_dim], + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.in_proj_z"), + ) + self.in_proj_b = LinearColParallelMerged( + config.hidden_size, + [config.linear_num_value_heads], + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.in_proj_b"), + ) + self.in_proj_a = LinearColParallelMerged( + config.hidden_size, + [config.linear_num_value_heads], + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.in_proj_a"), + ) + self.conv1d = _Qwen3_5DepthwiseConv1d( + local_conv_dim, config.linear_conv_kernel_dim + ) + self.dt_bias = torch.empty(local_value_heads) + self.A_log = torch.empty(local_value_heads, dtype=torch.float32) + self.norm = RMSNorm(config.linear_value_head_dim, eps=config.rms_norm_eps) + self.out_proj = LinearOProj( + value_dim, + config.hidden_size, + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.out_proj"), + ) + self._num_key_heads = local_key_heads + self._num_value_heads = local_value_heads + self._key_head_dim = config.linear_key_head_dim + self._value_head_dim = config.linear_value_head_dim + self._local_key_dim = local_key_heads * self._key_head_dim + self._local_value_dim = local_value_heads * self._value_head_dim + self._norm_eps = config.rms_norm_eps + self._layer_id = layer_id + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + ctx = get_global_ctx() + batch = ctx.batch + state_pool = ctx.gdn_state_pool + if state_pool is None: + raise RuntimeError("Qwen3.5 GDN state pool is not initialized") + layer_state = state_pool.layer_state(self._layer_id) + + mixed_qkv = self.in_proj_qkv.forward(hidden_states) + z = self.in_proj_z.forward(hidden_states).view( + -1, self._num_value_heads, self._value_head_dim + ) + b = self.in_proj_b.forward(hidden_states).view(-1, self._num_value_heads) + a = self.in_proj_a.forward(hidden_states).view(-1, self._num_value_heads) + + if batch.is_decode: + mixed_qkv = qwen35_causal_conv1d_decode( + mixed_qkv, + self.conv1d.weight, + layer_state.conv, + batch.state_indices, + batch.positions, + ) + else: + cu_seqlens, sequence_slots, first_positions = self._prefill_state_plan(batch) + initial_conv = layer_state.conv.index_select(0, sequence_slots) + mixed_qkv, final_conv = qwen35_causal_conv1d_prefill( + mixed_qkv, + self.conv1d.weight, + initial_conv, + cu_seqlens, + first_positions, + ) + layer_state.conv.index_copy_(0, sequence_slots, final_conv) + + query, key, value = mixed_qkv.split( + [self._local_key_dim, self._local_key_dim, self._local_value_dim], + dim=-1, + ) + query = query.view(-1, self._num_key_heads, self._key_head_dim) + key = key.view(-1, self._num_key_heads, self._key_head_dim) + value = value.view(-1, self._num_value_heads, self._value_head_dim) + + if batch.is_decode: + core_output = qwen35_gdn_decode( + query, + key, + value, + a, + b, + self.A_log, + self.dt_bias, + layer_state.recurrent, + batch.state_indices, + batch.positions, + ) + else: + g, beta = qwen35_gdn_gates(a, b, self.A_log, self.dt_bias) + initial_recurrent = layer_state.recurrent.index_select(0, sequence_slots) + keep = (first_positions != 0).to(initial_recurrent.dtype) + initial_recurrent = initial_recurrent * keep[:, None, None, None] + core_output, final_recurrent = qwen35_gdn_prefill( + query, + key, + value, + g, + beta, + initial_recurrent, + cu_seqlens, + ) + layer_state.recurrent.index_copy_(0, sequence_slots, final_recurrent) + + gated = qwen35_rmsnorm_gated( + core_output, z, self.norm.weight, self._norm_eps + ) + return self.out_proj.forward(gated.flatten(1)) + + @staticmethod + def _prefill_state_plan(batch): + cached = getattr(batch, "_qwen35_gdn_prefill_plan", None) + if cached is not None: + return cached + lengths = [int(req.extend_len) for req in batch.padded_reqs] + boundaries = [0] + for length in lengths: + boundaries.append(boundaries[-1] + length) + cu_seqlens = torch.tensor( + boundaries, dtype=torch.int32, device=batch.positions.device + ) + starts = cu_seqlens[:-1].to(torch.int64) + sequence_slots = batch.state_indices.index_select(0, starts) + first_positions = batch.positions.index_select(0, starts) + plan = (cu_seqlens, sequence_slots, first_positions) + batch._qwen35_gdn_prefill_plan = plan + return plan + + +class Qwen3_5FullAttention(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + tp_size = get_tp_info().size + self.num_qo_heads = div_even(config.num_qo_heads, tp_size) + self.num_kv_heads = div_even(config.num_kv_heads, tp_size, allow_replicate=True) + self.head_dim = config.head_dim + self.q_size = self.num_qo_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + module_prefix = f"model.layers.{layer_id}.self_attn" + self.qkv_proj = LinearQKVGatedMerged( + config.hidden_size, + config.head_dim, + config.num_qo_heads, + config.num_kv_heads, + has_bias=False, + quant=config.quant_for_modules( + f"{module_prefix}.q_proj", + f"{module_prefix}.k_proj", + f"{module_prefix}.v_proj", + ), + ) + self.q_norm = RMSNormOffset(config.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNormOffset(config.head_dim, eps=config.rms_norm_eps) + self.attn = AttentionLayer( + layer_id=layer_id, + head_dim=config.head_dim, + num_qo_heads=config.num_qo_heads, + num_kv_heads=config.num_kv_heads, + rotary_config=config.rotary_config, + q_norm=self.q_norm, + k_norm=self.k_norm, + ) + self.o_proj = LinearOProj( + config.num_qo_heads * config.head_dim, + config.hidden_size, + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.o_proj"), + ) + + def split_query_gate(self, q_gate: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + leading = q_gate.shape[:-1] + per_head = q_gate.view(*leading, self.num_qo_heads, 2 * self.head_dim) + query, gate = per_head.chunk(2, dim=-1) + return query.reshape(*leading, self.q_size), gate.reshape(*leading, self.q_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + qkv = self.qkv_proj.forward(hidden_states) + q_gate, key, value = qkv.split( + [2 * self.q_size, self.kv_size, self.kv_size], dim=-1 + ) + query, gate = self.split_query_gate(q_gate) + attn_output = self.attn.forward(torch.cat((query, key, value), dim=-1)) + return self.o_proj.forward(attn_output * torch.sigmoid(gate)) + + +class Qwen3_5SharedExpert(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + intermediate_size = config.shared_expert_intermediate_size + module_prefix = f"model.layers.{layer_id}.mlp.shared_expert" + self.gate_up_proj = LinearColParallelMerged( + config.hidden_size, + [intermediate_size, intermediate_size], + has_bias=False, + quant=config.quant_for_modules( + f"{module_prefix}.gate_proj", + f"{module_prefix}.up_proj", + ), + ) + self.down_proj = LinearRowParallel( + intermediate_size, + config.hidden_size, + has_bias=False, + quant=config.quant_for_modules(f"{module_prefix}.down_proj"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(hidden_states))) + + +def qwen35_shared_output( + shared_hidden: torch.Tensor, gate_logits: torch.Tensor +) -> torch.Tensor: + return shared_hidden * torch.sigmoid(gate_logits) + + +class Qwen3_5MoeMLP(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.experts = MoELayer( + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant=config.quant, + ) + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.shared_expert = Qwen3_5SharedExpert(config, layer_id) + self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + routed = self.experts.forward( + hidden_states=hidden_states, + router_logits=self.gate.forward(hidden_states), + ) + shared = qwen35_shared_output( + self.shared_expert.forward(hidden_states), + self.shared_expert_gate.forward(hidden_states), + ) + return routed + shared + + +class Qwen3_5DecoderLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.layer_type = config.layer_types[layer_id] + if self.layer_type == "linear_attention": + self.linear_attn = Qwen3_5LinearAttention(config, layer_id) + elif self.layer_type == "full_attention": + self.self_attn = Qwen3_5FullAttention(config, layer_id) + else: + raise ValueError(f"Invalid Qwen3.5 layer type: {self.layer_type}") + self.mlp = Qwen3_5MoeMLP(config, layer_id) + self.input_layernorm = RMSNormOffsetFused(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = RMSNormOffsetFused( + config.hidden_size, config.rms_norm_eps + ) + self._layer_id = layer_id + + def forward_token_mixer(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.layer_type == "linear_attention": + return self.linear_attn.forward(hidden_states) + return self.self_attn.forward(hidden_states) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward( + self, hidden_states: torch.Tensor, residual: torch.Tensor | None = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.input_layernorm.forward(hidden_states, residual) + hidden_states = self.forward_token_mixer(hidden_states) + hidden_states, residual = self.post_attention_layernorm.forward( + hidden_states, residual + ) + return self.mlp.forward(hidden_states), residual + + +class Qwen3_5Model(BaseOP): + def __init__(self, config: ModelConfig): + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = OPList( + [Qwen3_5DecoderLayer(config, layer_id) for layer_id in range(config.num_layers)] + ) + self.norm = RMSNormOffsetFused(config.hidden_size, config.rms_norm_eps) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + hidden_states = self.embed_tokens.forward(input_ids) + residual = None + for layer in self.layers.op_list: + hidden_states, residual = layer.forward(hidden_states, residual) + return self.norm.forward(hidden_states, residual)[0] + + +class Qwen3_5MoeForConditionalGeneration(BaseLLMModel): + """Text-only Qwen3.5-MoE implementation for the multimodal checkpoint.""" + + def __init__(self, config: ModelConfig): + self.model = Qwen3_5Model(config) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + ) + super().__init__() + + def forward(self) -> torch.Tensor: + hidden_states = self.model.forward(get_global_ctx().batch.input_ids) + return self.lm_head.forward(hidden_states) + + +__all__ = [ + "Qwen3_5FullAttention", + "Qwen3_5LinearAttention", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5SharedExpert", + "qwen35_shared_output", +] diff --git a/python/minisgl/models/qwen35_moe_afd.py b/python/minisgl/models/qwen35_moe_afd.py new file mode 100644 index 0000000..28d1796 --- /dev/null +++ b/python/minisgl/models/qwen35_moe_afd.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from typing import NamedTuple + +import torch +from minisgl.layers import ( + BaseOP, + LinearReplicated, + MoELayer, + OPList, + ParallelLMHead, + RMSNormOffsetFused, + VocabParallelEmbedding, +) +from minisgl.models.afd import ModelStageState +from minisgl.utils import nvtx_annotate + +from .config import ModelConfig +from .qwen35_moe import ( + Qwen3_5FullAttention, + Qwen3_5LinearAttention, + Qwen3_5SharedExpert, + qwen35_shared_output, +) +from .qwen3_moe_afd import qwen3_afd_topk + + +class Qwen3_5AfdTopK(NamedTuple): + topk_ids: torch.Tensor | None + topk_weights: torch.Tensor | None + router_logits: torch.Tensor | None = None + renormalize: bool = True + deepep_topk: bool = False + shared_output: torch.Tensor | None = None + dispatch_fp8: bool = False + + +def _qwen35_dispatch_uses_fp8(config: ModelConfig) -> bool: + return bool(config.quant is not None and config.quant.method.startswith("fp8")) + + +class _Qwen3_5AfdGateAndShared(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.shared_expert = Qwen3_5SharedExpert(config, layer_id) + self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False) + + +class Qwen3_5AfdAGLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.layer_type = config.layer_types[layer_id] + if self.layer_type == "linear_attention": + self.linear_attn = Qwen3_5LinearAttention(config, layer_id) + elif self.layer_type == "full_attention": + self.self_attn = Qwen3_5FullAttention(config, layer_id) + else: + raise ValueError(f"Invalid Qwen3.5 layer type: {self.layer_type}") + self.mlp = _Qwen3_5AfdGateAndShared(config, layer_id) + self.input_layernorm = RMSNormOffsetFused(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = RMSNormOffsetFused( + config.hidden_size, config.rms_norm_eps + ) + self.top_k = config.num_experts_per_tok + self.renormalize = config.norm_topk_prob + self.dispatch_fp8 = _qwen35_dispatch_uses_fp8(config) + self._layer_id = layer_id + + @nvtx_annotate("AFD_AG_Attn_{}", layer_id_field="_layer_id") + def forward_attention( + self, hidden_states: torch.Tensor, residual: torch.Tensor | None = None + ) -> ModelStageState: + hidden_states, residual = self.input_layernorm.forward(hidden_states, residual) + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn.forward(hidden_states) + else: + hidden_states = self.self_attn.forward(hidden_states) + hidden_states, residual = self.post_attention_layernorm.forward( + hidden_states, residual + ) + return ModelStageState(hidden_states=hidden_states, residual=residual) + + @nvtx_annotate("AFD_AG_Route_{}", layer_id_field="_layer_id") + def route(self, hidden_states: torch.Tensor) -> Qwen3_5AfdTopK: + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + base = qwen3_afd_topk( + self.mlp.gate.forward(hidden_states), + top_k=self.top_k, + renormalize=self.renormalize, + dispatch_fp8=self.dispatch_fp8, + ) + shared = qwen35_shared_output( + self.mlp.shared_expert.forward(hidden_states), + self.mlp.shared_expert_gate.forward(hidden_states), + ) + return Qwen3_5AfdTopK( + topk_ids=base.topk_ids, + topk_weights=base.topk_weights, + router_logits=base.router_logits, + renormalize=base.renormalize, + deepep_topk=base.deepep_topk, + shared_output=shared, + dispatch_fp8=self.dispatch_fp8, + ) + + +class Qwen3_5AfdDenseRouterStage(BaseOP): + def __init__(self, config: ModelConfig): + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = OPList( + [Qwen3_5AfdAGLayer(config, layer_id) for layer_id in range(config.num_layers)] + ) + self.norm = RMSNormOffsetFused(config.hidden_size, config.rms_norm_eps) + + def embed_input_ids(self, input_ids: torch.Tensor) -> ModelStageState: + return ModelStageState(self.embed_tokens.forward(input_ids), residual=None) + + def forward_attention(self, layer_id: int, hidden_states, residual=None): + return self.layers.op_list[layer_id].forward_attention(hidden_states, residual) + + def route(self, layer_id: int, hidden_states): + return self.layers.op_list[layer_id].route(hidden_states) + + def finalize_hidden(self, hidden_states, residual=None): + return self.norm.forward(hidden_states, residual)[0] + + +class Qwen3_5AfdDenseRouterForCausalLM(BaseOP): + def __init__(self, config: ModelConfig): + self.model = Qwen3_5AfdDenseRouterStage(config) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + ) + + def embed_input_ids(self, input_ids): + return self.model.embed_input_ids(input_ids) + + def forward_attention(self, layer_id, hidden_states, residual=None): + return self.model.forward_attention(layer_id, hidden_states, residual) + + def route(self, layer_id, hidden_states): + return self.model.route(layer_id, hidden_states) + + def finalize_hidden(self, hidden_states, residual=None): + return self.model.finalize_hidden(hidden_states, residual) + + def forward_lm_head(self, hidden_states): + return self.lm_head.forward(hidden_states) + + +class _Qwen3_5AfdExpertMLP(BaseOP): + def __init__(self, config: ModelConfig): + self.experts = MoELayer( + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant=config.quant, + a2a_backend="none", + ) + + +class Qwen3_5AfdEGLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.mlp = _Qwen3_5AfdExpertMLP(config) + self._layer_id = layer_id + + def run_experts(self, dispatch_output) -> torch.Tensor: + return self.mlp.experts._run_moe_core(dispatch_output).hidden_states + + +class _Qwen3_5AfdExpertStageInner(BaseOP): + def __init__(self, config: ModelConfig): + self.layers = OPList( + [Qwen3_5AfdEGLayer(config, layer_id) for layer_id in range(config.num_layers)] + ) + + def run_experts(self, layer_id: int, dispatch_output) -> torch.Tensor: + return self.layers.op_list[layer_id].run_experts(dispatch_output) + + +class Qwen3_5AfdExpertStage(BaseOP): + def __init__(self, config: ModelConfig): + self.model = _Qwen3_5AfdExpertStageInner(config) + self._dispatch_fp8 = _qwen35_dispatch_uses_fp8(config) + + def dispatch_uses_fp8(self) -> bool: + return self._dispatch_fp8 + + def run_experts(self, layer_id: int, dispatch_output) -> torch.Tensor: + return self.model.run_experts(layer_id, dispatch_output) + + +__all__ = [ + "Qwen3_5AfdAGLayer", + "Qwen3_5AfdDenseRouterForCausalLM", + "Qwen3_5AfdExpertStage", + "Qwen3_5AfdTopK", +] diff --git a/python/minisgl/models/qwen3_moe_afd.py b/python/minisgl/models/qwen3_moe_afd.py index 5eff365..268680e 100644 --- a/python/minisgl/models/qwen3_moe_afd.py +++ b/python/minisgl/models/qwen3_moe_afd.py @@ -84,7 +84,7 @@ def qwen3_afd_topk( dispatch_fp8=dispatch_fp8, ) - from minisgl.kernel import topk_gating + from minisgl.kernel import topk_softmax num_tokens = router_logits.shape[0] weights = torch.empty( @@ -97,11 +97,10 @@ def qwen3_afd_topk( dtype=torch.int32, device=router_logits.device, ) - topk_gating( + topk_softmax( weights, ids, router_logits.contiguous(), - scoring="softmax", renormalize=renormalize, ) return Qwen3AfdTopK( diff --git a/python/minisgl/models/register.py b/python/minisgl/models/register.py index 423895f..66b0551 100644 --- a/python/minisgl/models/register.py +++ b/python/minisgl/models/register.py @@ -7,6 +7,10 @@ "Qwen2ForCausalLM": (".qwen2", "Qwen2ForCausalLM"), "Qwen3ForCausalLM": (".qwen3", "Qwen3ForCausalLM"), "Qwen3MoeForCausalLM": (".qwen3_moe", "Qwen3MoeForCausalLM"), + "Qwen3_5MoeForConditionalGeneration": ( + ".qwen35_moe", + "Qwen3_5MoeForConditionalGeneration", + ), "MiniMaxM2ForCausalLM": (".minimax_m2", "MiniMaxM2ForCausalLM"), "Glm4MoeForCausalLM": (".glm4_moe", "Glm4MoeForCausalLM"), "MistralForCausalLM": (".mistral", "MistralForCausalLM"), diff --git a/python/minisgl/models/weight.py b/python/minisgl/models/weight.py index 46b984f..fb9116d 100644 --- a/python/minisgl/models/weight.py +++ b/python/minisgl/models/weight.py @@ -19,6 +19,7 @@ _FP8_SCALE_SUFFIX = ".weight_scale_inv" _FP8_SCALE_SUFFIXES = (".weight_scale_inv", ".weight_scale") _SCALE_RENAME_SUFFIX = "_scale" +_NVFP4_GLOBAL_SCALE_SUFFIXES = (".input_scale", ".weight_scale_2") # Merge groups: individual projections -> fused projection _MERGE_GROUPS = { @@ -64,6 +65,39 @@ def _is_glm4_moe_config(config: Any) -> bool: ) == "glm4_moe" +def _is_qwen35_config(config: Any) -> bool: + archs = _config_field(config, "architectures", default=[]) or [] + return ( + "Qwen3_5MoeForConditionalGeneration" in archs + or _config_field(config, "model_type", default="") + in {"qwen3_5", "qwen3_5_moe", "qwen3_5_moe_text"} + ) + + +def _normalize_checkpoint_key(raw_name: str, config: Any) -> str | None: + """Map an HF checkpoint key to mini-sgl's text-model namespace. + + Qwen3.5 wraps the language model below ``model.language_model``. Vision, + MTP, and calibrated FP8-KV scale tensors are intentionally outside the + current text/BF16-KV runtime contract. + """ + + if raw_name.startswith(("vision_tower.", "multi_modal_projector.")): + return None + if _is_qwen35_config(config): + if raw_name.startswith(("model.visual.", "visual.", "mtp.", "model.mtp.")): + return None + if raw_name.endswith((".k_proj.k_scale", ".v_proj.v_scale")): + return None + if raw_name.startswith("model.language_model."): + raw_name = "model." + raw_name.removeprefix("model.language_model.") + elif raw_name.startswith("language_model."): + raw_name = raw_name.removeprefix("language_model.") + else: + raw_name = raw_name.removeprefix("language_model.") + return _rename_checkpoint_key(raw_name) + + def _glm4_layer_index_for_key(key: str) -> int | None: match = _GLM4_LAYER_PATTERN.match(key.removeprefix("language_model.")) return int(match.group("idx")) if match is not None else None @@ -121,6 +155,92 @@ def _shard_tensor( return value +def _shard_qwen35_linear_tensor( + key: str, + value: torch.Tensor, + r: int, + n: int, + config: Any, + *, + fp8_block_size: Optional[Tuple[int, int]] = None, +) -> torch.Tensor | None: + """Shard Qwen3.5 GDN tensors whose names do not match generic QKV rules.""" + + marker = ".linear_attn." + if marker not in key: + return None + key_dim = int(config.linear_num_key_heads) * int(config.linear_key_head_dim) + value_dim = int(config.linear_num_value_heads) * int( + config.linear_value_head_dim + ) + + if key.endswith(".in_proj_qkv.weight_scale"): + if fp8_block_size is None: + return None + block_out = int(fp8_block_size[0]) + if key_dim % block_out or value_dim % block_out: + raise ValueError( + "Qwen3.5 fused GDN FP8 scale requires block-aligned logical " + f"projections, got key={key_dim}, value={value_dim}, block={block_out}" + ) + key_blocks = key_dim // block_out + value_blocks = value_dim // block_out + q, k, v = value.split((key_blocks, key_blocks, value_blocks), dim=0) + return torch.cat( + ( + q.chunk(n, dim=0)[r], + k.chunk(n, dim=0)[r], + v.chunk(n, dim=0)[r], + ), + dim=0, + ).clone() + if key.endswith( + ( + ".in_proj_z.weight_scale", + ".in_proj_a.weight_scale", + ".in_proj_b.weight_scale", + ) + ): + if fp8_block_size is None: + return None + return _slice_fp8_scale_dim( + value, + dim=0, + r=r, + n=n, + block_size=int(fp8_block_size[0]), + ) + if key.endswith(".out_proj.weight_scale"): + if fp8_block_size is None: + return None + return _slice_fp8_scale_dim( + value, + dim=1, + r=r, + n=n, + block_size=int(fp8_block_size[1]), + ) + if key.endswith((".in_proj_qkv.weight", ".conv1d.weight")): + q, k, v = value.split((key_dim, key_dim, value_dim), dim=0) + return torch.cat( + (q.chunk(n, dim=0)[r], k.chunk(n, dim=0)[r], v.chunk(n, dim=0)[r]), + dim=0, + ).clone() + if key.endswith( + ( + ".in_proj_z.weight", + ".in_proj_a.weight", + ".in_proj_b.weight", + ".A_log", + ".dt_bias", + ) + ): + return value.chunk(n, dim=0)[r].clone() + if key.endswith(".out_proj.weight"): + return value.chunk(n, dim=1)[r].clone() + return None + + def _slice_fp8_scale_dim( value: torch.Tensor, *, @@ -187,6 +307,55 @@ def _shard_fp8_channel_scale_tensor( return value +def _shard_nvfp4_tensor( + key: str, + value: torch.Tensor, + r: int, + n: int, + num_kv_heads: int, + *, + group_size: int, + skip_tp_shard: bool = False, +) -> torch.Tensor: + """Shard a ModelOpt NVFP4 tensor in its serialized packed layout.""" + + if skip_tp_shard: + return value + if key.endswith(_NVFP4_GLOBAL_SCALE_SUFFIXES): + return value + if key.endswith(".weight_scale"): + return _shard_fp8_scale_tensor( + key, + value, + r, + n, + num_kv_heads, + (1, int(group_size)), + ) + return _shard_tensor(key, value, r, n, num_kv_heads) + + +def _nvfp4_runtime_parameter_name(name: str) -> str: + """Map ModelOpt projection suffixes to flat MoELayer tensor attributes.""" + + for checkpoint_suffix, runtime_suffix in ( + (".weight_scale_2", "_weight_scale_2"), + (".input_scale", "_input_scale"), + (".weight_scale", "_scale"), + ): + if name.endswith(checkpoint_suffix): + return name.removesuffix(checkpoint_suffix) + runtime_suffix + return name + + +def _merge_projection_parts(parts: list[torch.Tensor]) -> torch.Tensor: + """Merge projection tensors, including scalar ModelOpt global scales.""" + + if not parts: + raise ValueError("cannot merge an empty projection group") + return torch.stack(parts, dim=0) if parts[0].dim() == 0 else torch.cat(parts, dim=0) + + def _find_paired_scale_key( *, normalized_weight_name: str, @@ -308,23 +477,26 @@ def load_weight( else None ) fp8_channel_scale = config.quant is not None and config.quant.method == "fp8_channel" + is_nvfp4 = config.quant is not None and config.quant.method == "modelopt_nvfp4" + nvfp4_group_size = int(config.quant.group_size or 16) if is_nvfp4 else 0 logger.info( f"Loading weights from {len(files)} safetensors files" + (f" (FP8 block={fp8_block_size})" if fp8_block_size else "") + (" (FP8 channel)" if fp8_channel_scale else "") + + (f" (ModelOpt NVFP4 group={nvfp4_group_size})" if is_nvfp4 else "") ) # Buffer for merge groups: merged_key -> {slot: tensor} merge_buf: Dict[str, Dict[str, torch.Tensor]] = {} expert_buf: Dict[str, Dict[int, torch.Tensor]] = {} + nvfp4_input_scale_max: Dict[str, torch.Tensor] = {} for file in tqdm(files, desc="Loading weights", disable=not tp_info.is_primary()): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: shard_keys = set(f.keys()) for raw_name in f.keys(): - # Strip multimodal wrapper prefix, skip vision/projector weights - if raw_name.startswith(("vision_tower.", "multi_modal_projector.")): + name = _normalize_checkpoint_key(raw_name, config) + if name is None: continue - name = raw_name.removeprefix("language_model.") if _is_glm4_moe_config(config) and _glm4_is_expected_nextn_key(name, config): continue @@ -343,18 +515,26 @@ def load_weight( if is_scale_tensor: assert scale_suffix is not None base_weight_name = name.removesuffix(scale_suffix) - if fp8_block_size is None and not fp8_channel_scale: + if fp8_block_size is None and not fp8_channel_scale and not is_nvfp4: continue - name = ( - base_weight_name + _SCALE_RENAME_SUFFIX - if _is_expert_key(base_weight_name) - else base_weight_name + ".weight_scale" - ) - name = _rename_checkpoint_key(name) - + if is_nvfp4 and _is_expert_key(base_weight_name): + name = _nvfp4_runtime_parameter_name(name) + else: + name = ( + base_weight_name + _SCALE_RENAME_SUFFIX + if _is_expert_key(base_weight_name) + else base_weight_name + ".weight_scale" + ) pre_merge_expert_info = ( _get_expert_stack_info(name) if config.is_moe else None ) + is_nvfp4_input_scale = bool( + is_nvfp4 + and pre_merge_expert_info is not None + and name.endswith(".input_scale") + ) + if is_nvfp4 and pre_merge_expert_info is not None: + name = _nvfp4_runtime_parameter_name(name) is_remote_shared_expert = _is_remote_shared_expert_key(name, config) # afd AG/EG role split: drop the other role's weights. if skip_expert_weights and ( @@ -365,7 +545,11 @@ def load_weight( pre_merge_expert_info is None and not is_remote_shared_expert ): continue - if local_expert_partition is not None and pre_merge_expert_info is not None: + if ( + local_expert_partition is not None + and pre_merge_expert_info is not None + and not is_nvfp4_input_scale + ): _, expert_idx = pre_merge_expert_info if not ( local_expert_partition.start_idx @@ -374,7 +558,44 @@ def load_weight( ): continue raw = f.get_tensor(raw_name) - if fp8_channel_scale and is_scale_tensor: + if is_nvfp4_input_scale: + assert pre_merge_expert_info is not None + packed_input_key, _ = pre_merge_expert_info + packed_input_key = _nvfp4_runtime_parameter_name(packed_input_key) + input_scale = raw.float() + previous = nvfp4_input_scale_max.get(packed_input_key) + nvfp4_input_scale_max[packed_input_key] = ( + input_scale + if previous is None + else torch.maximum(previous, input_scale) + ) + del raw + continue + qwen35_linear_tensor = ( + _shard_qwen35_linear_tensor( + name, + raw, + tp_info.rank, + tp_info.size, + config, + fp8_block_size=fp8_block_size, + ) + if _is_qwen35_config(config) + else None + ) + if qwen35_linear_tensor is not None: + tensor = qwen35_linear_tensor + elif is_nvfp4 and pre_merge_expert_info is not None: + tensor = _shard_nvfp4_tensor( + _normalize_checkpoint_key(raw_name, config) or name, + raw, + tp_info.rank, + tp_info.size, + config.num_kv_heads, + group_size=nvfp4_group_size, + skip_tp_shard=(local_expert_partition is not None), + ) + elif fp8_channel_scale and is_scale_tensor: tensor = _shard_fp8_channel_scale_tensor( name, raw, @@ -431,7 +652,7 @@ def load_weight( continue parts = [merge_buf[merged_key][s] for s in all_slots] del merge_buf[merged_key] - out = (merged_key, torch.cat(parts, dim=0)) + out = (merged_key, _merge_projection_parts(parts)) if config.is_moe and (expert_info := _get_expert_stack_info(out[0])) is not None: packed_key, expert_idx = expert_info @@ -455,12 +676,27 @@ def load_weight( experts = [slots.get(idx, zero) for idx in range(local_num_experts)] del expert_buf[packed_key] stacked = torch.stack(experts, dim=0) - if packed_key.endswith(_SCALE_RENAME_SUFFIX): + if packed_key.endswith(_SCALE_RENAME_SUFFIX) and not is_nvfp4: stacked = stacked.to(torch.float32) yield packed_key, stacked else: # Normal dense model yield out[0], out[1] + nvfp4_input_merge: Dict[str, Dict[str, torch.Tensor]] = {} + for key, scale in nvfp4_input_scale_max.items(): + if (info := _get_merge_info(key)) is None: + yield key, scale + continue + merged_key, slot, all_slots = info + slots = nvfp4_input_merge.setdefault(merged_key, {}) + slots[slot] = scale + if all(name in slots for name in all_slots): + yield merged_key, torch.stack([slots[name] for name in all_slots]) + del nvfp4_input_merge[merged_key] + assert not merge_buf, f"Incomplete merge groups in checkpoint: {list(merge_buf.keys())}" assert not expert_buf, f"Incomplete expert tensors in checkpoint: {list(expert_buf.keys())}" + assert not nvfp4_input_merge, ( + f"Incomplete NVFP4 input-scale merge groups: {list(nvfp4_input_merge.keys())}" + ) logger.info("Finished loading weights") diff --git a/python/minisgl/moe/deepep_m2n_adapter.py b/python/minisgl/moe/deepep_m2n_adapter.py index b8e5c79..18afbc4 100644 --- a/python/minisgl/moe/deepep_m2n_adapter.py +++ b/python/minisgl/moe/deepep_m2n_adapter.py @@ -50,6 +50,7 @@ def __init__( hidden_size: int, top_k: int, num_max_dispatch_tokens_per_rank: int, + use_expanded_layout: bool = True, log: Callable[[str], None] | None = None, ) -> None: if not dist.is_initialized(): @@ -75,6 +76,7 @@ def __init__( self.is_eg = not self.is_ag self.hidden_size = int(hidden_size) self.top_k = int(top_k) + self.use_expanded_layout = bool(use_expanded_layout) self.num_max_dispatch_tokens_per_rank = int(num_max_dispatch_tokens_per_rank) self._deepep_expert_map_by_device: dict[tuple[str, int | None, str], torch.Tensor] = {} if self.hidden_size <= 0: @@ -183,9 +185,14 @@ def dispatch( deepep_topk_ids.contiguous(), real_topk_weights.float().contiguous(), hidden_states_scale=hidden_states_scale, - expert_alignment=self._resolve_expert_alignment(expert_alignment), + expert_alignment=( + self._resolve_expert_alignment(expert_alignment) + if self.use_expanded_layout + else 1 + ), num_max_dispatch_tokens_per_rank=effective_max_tokens, do_cpu_sync=bool(do_cpu_sync), + use_expanded_layout=self.use_expanded_layout, ) if recv_topk_weights is not None and int(recv_topk_weights.shape[0]) > int(recv_x.shape[0]): recv_topk_weights = recv_topk_weights[: recv_x.shape[0]] @@ -216,7 +223,7 @@ def combine( if topk_weights is not None: raise RuntimeError( "DeepEP M2N combine must not receive topk_weights; " - "the M2N DeepGEMM path applies expanded top-k weights before combine" + "the expert runner applies routed top-k weights before combine" ) handle = dispatch_output.handle else: diff --git a/python/minisgl/scheduler/scheduler.py b/python/minisgl/scheduler/scheduler.py index 506a064..8ee1d20 100644 --- a/python/minisgl/scheduler/scheduler.py +++ b/python/minisgl/scheduler/scheduler.py @@ -140,6 +140,7 @@ def _materialize_batch_metadata( batch.positions = positions_device batch.positions_host = positions_host input_mapping = _make_input_tuple(batch, device) + batch.state_indices = input_mapping[0] write_mapping = _make_write_tuple(batch, device) batch.out_loc = page_table[input_mapping] if token_pool is not None: @@ -151,6 +152,11 @@ class Scheduler(SchedulerIOMixin): def __init__(self, config: SchedulerConfig): from minisgl.engine import Engine + if config.model_config.is_hybrid_attention and config.cache_type != "naive": + raise RuntimeError( + "Hybrid GDN models require cache_type='naive' until GDN prefix " + "state snapshots are implemented" + ) self.engine = Engine(config) # use another stream to overlap metadata processing with computation diff --git a/python/minisgl/server/args.py b/python/minisgl/server/args.py index 1cf5fc4..f22e287 100644 --- a/python/minisgl/server/args.py +++ b/python/minisgl/server/args.py @@ -12,6 +12,23 @@ from minisgl.utils import init_logger +def _hf_config_uses_modelopt_nvfp4(config) -> bool: + def value(obj, name, default=None): + return obj.get(name, default) if isinstance(obj, dict) else getattr(obj, name, default) + + for candidate in (config, value(config, "text_config", None)): + if candidate is None: + continue + quant = value(candidate, "quantization_config", None) + if quant is None: + continue + if not isinstance(quant, dict): + quant = getattr(quant, "to_dict", lambda: {})() + if quant.get("quant_method") == "modelopt" and quant.get("quant_algo") == "NVFP4": + return True + return False + + @dataclass(frozen=True) class ServerArgs(SchedulerConfig): mode: str = "serve" @@ -47,6 +64,8 @@ class ServerArgs(SchedulerConfig): afd_report_dir: str = "" afd_device_comm_num_sms: int = 1 afd_disable_overlap: bool = False + afd_async_pending_steps: int = 3 + afd_force_multi_mb_graph_overlap: bool = False afd_num_mb: int = 1 @property @@ -178,6 +197,33 @@ def _validate_afd_parallel_args( ) +def _afd_multi_mb_graph_requires_serial_steps(kwargs: dict) -> bool: + """DeepEP graph instances are not safe with multiple steps in flight. + + This only serializes the coordinator/worker command lead. The per-step + multi-MB AG/EG paper pipeline still overlaps its compute and communication + lanes. + """ + graph_requested = bool(kwargs.get("afd_decode_graph_bs")) or int( + kwargs.get("cuda_graph_max_bs") or 0 + ) > 0 + return ( + kwargs.get("mode") == "afd-serve" + and int(kwargs.get("afd_num_mb") or 1) > 1 + and bool(kwargs.get("afd_enable_attention_decode_graph")) + and bool(kwargs.get("afd_enable_model_decode_graph")) + and graph_requested + ) + + +def _afd_multi_mb_graph_auto_serializes(kwargs: dict) -> bool: + return ( + _afd_multi_mb_graph_requires_serial_steps(kwargs) + and not bool(kwargs.get("afd_disable_overlap")) + and not bool(kwargs.get("afd_force_multi_mb_graph_overlap")) + ) + + def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bool]: """ Parse command line arguments and return an EngineConfig. @@ -511,12 +557,21 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo parser.add_argument( "--afd-moe-runner-backend", - choices=["auto", "triton", "triton_fp8", "triton_kernel", "deep_gemm"], + choices=[ + "auto", + "triton", + "triton_fp8", + "triton_kernel", + "deep_gemm", + "flashinfer_nvfp4", + ], default=ServerArgs.afd_moe_runner_backend, help=( "MoE compute runner backend. 'triton_fp8' uses the FP8-native " "Triton fallback; 'deep_gemm' uses minisgl's vendored grouped " - "DeepGEMM kernels for DeepEP V2 expanded layout." + "DeepGEMM kernels for DeepEP V2 expanded layout; " + "'flashinfer_nvfp4' uses native ModelOpt NVFP4 routed experts with " + "DeepEP non-expanded token/top-k transport." ), ) @@ -629,6 +684,26 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo help="Disable AFD overlap scheduling on both attention and model workers (overlap is on by default).", ) + parser.add_argument( + "--afd-async-pending-steps", + type=int, + default=ServerArgs.afd_async_pending_steps, + help=( + "Maximum number of outstanding global AFD steps in overlap mode. " + "The workers launch this window before retiring the oldest reply." + ), + ) + + parser.add_argument( + "--afd-force-multi-mb-graph-overlap", + action="store_true", + default=ServerArgs.afd_force_multi_mb_graph_overlap, + help=( + "Diagnostic opt-in: allow cross-step overlap for multi-MB decode " + "graphs despite the collective re-entry safety guard." + ), + ) + parser.add_argument( "--afd-num-mb", type=int, @@ -664,10 +739,12 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo kwargs["model_path"] = model_path del kwargs["model_source"] + hf_config = None if (dtype_str := kwargs["dtype"]) == "auto": from minisgl.utils import cached_load_hf_config - dtype_str = cached_load_hf_config(kwargs["model_path"]).dtype + hf_config = cached_load_hf_config(kwargs["model_path"]) + dtype_str = hf_config.dtype DTYPE_MAP = { "float16": torch.float16, @@ -688,12 +765,19 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo parser.error( "AFD serve centralized scheduler currently supports only --cache-type naive" ) - # AFD experts run DeepEP m2n dispatch/combine under a DeepGEMM contract; the - # generic auto->triton runner default deadlocks the cross-role NVLink barrier - # (Triton experts vs DeepGEMM-contract comm). Default the AFD expert runner to - # deep_gemm -- the only backend the AFD EG path supports. if kwargs["afd_moe_runner_backend"] == "auto": - kwargs["afd_moe_runner_backend"] = "deep_gemm" + if hf_config is None: + from minisgl.utils import cached_load_hf_config + + hf_config = cached_load_hf_config(kwargs["model_path"]) + # ModelOpt NVFP4 checkpoints need FlashInfer's native W4A4 MoE + # runner and DeepEP's non-expanded token/top-k transport. Existing + # BF16/FP8 AFD models retain the expanded DeepGEMM contract. + kwargs["afd_moe_runner_backend"] = ( + "flashinfer_nvfp4" + if _hf_config_uses_modelopt_nvfp4(hf_config) + else "deep_gemm" + ) if kwargs["cuda_graph_max_bs"] is None and not kwargs["afd_decode_graph_bs"]: cap = ( int(kwargs["afd_max_running_req"]) @@ -702,9 +786,35 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo ) if cap >= 1: kwargs["cuda_graph_max_bs"] = min(cap, 512) + if int(kwargs["afd_async_pending_steps"]) < 2: + parser.error("--afd-async-pending-steps must be >= 2") + if bool(kwargs["afd_force_multi_mb_graph_overlap"]) and bool( + kwargs["afd_disable_overlap"] + ): + parser.error( + "--afd-force-multi-mb-graph-overlap conflicts with --afd-disable-overlap" + ) + auto_serialized_graph_steps = _afd_multi_mb_graph_auto_serializes(kwargs) + if auto_serialized_graph_steps: + kwargs["afd_disable_overlap"] = True _validate_afd_parallel_args(parser, kwargs) result = ServerArgs(**kwargs) logger = init_logger(__name__) + if auto_serialized_graph_steps: + logger.warning( + "AFD multi-MB decode graphs serialize cross-step command launch to " + "avoid re-entering captured DeepEP collectives; within-step " + "microbatch compute/communication overlap remains enabled." + ) + elif ( + bool(kwargs["afd_force_multi_mb_graph_overlap"]) + and _afd_multi_mb_graph_requires_serial_steps(kwargs) + ): + logger.warning( + "AFD multi-MB decode graph cross-step overlap is FORCE ENABLED for " + "diagnostic profiling; async_pending_steps=%d", + int(kwargs["afd_async_pending_steps"]), + ) logger.info(f"Parsed arguments:\n{result}") return result, run_shell diff --git a/python/minisgl/server/supervisor_ray.py b/python/minisgl/server/supervisor_ray.py index fa9d410..32386f9 100644 --- a/python/minisgl/server/supervisor_ray.py +++ b/python/minisgl/server/supervisor_ray.py @@ -74,6 +74,9 @@ "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", diff --git a/python/minisgl/utils/hf.py b/python/minisgl/utils/hf.py index 5bd99de..abbd076 100644 --- a/python/minisgl/utils/hf.py +++ b/python/minisgl/utils/hf.py @@ -1,8 +1,9 @@ import functools +import json import os from typing import Any -from huggingface_hub import snapshot_download +from huggingface_hub import hf_hub_download, snapshot_download from transformers import AutoConfig, PretrainedConfig from minisgl.hf_support import DisabledTqdm, load_tokenizer, local_files_only, resolve_local_model_dir @@ -10,11 +11,32 @@ @functools.cache def _load_hf_config_cached(resolved_model_path: str, local_files_only: bool) -> Any: - return AutoConfig.from_pretrained( - resolved_model_path, - local_files_only=local_files_only, - trust_remote_code=True, - ) + try: + return AutoConfig.from_pretrained( + resolved_model_path, + local_files_only=local_files_only, + trust_remote_code=True, + ) + except ValueError: + # Transformers 4.x predates Qwen3.5, while vLLM ships a compatible + # config implementation. mini-sgl only needs the serialized fields, so + # retain them in a generic PretrainedConfig instead of requiring a + # process-wide Transformers upgrade. + if os.path.isdir(resolved_model_path): + config_path = os.path.join(resolved_model_path, "config.json") + if not os.path.isfile(config_path): + raise + else: + config_path = hf_hub_download( + repo_id=resolved_model_path, + filename="config.json", + local_files_only=local_files_only, + ) + with open(config_path, encoding="utf-8") as f: + config_dict = json.load(f) + if config_dict.get("model_type") not in {"qwen3_5", "qwen3_5_moe"}: + raise + return PretrainedConfig(**config_dict) def _load_hf_config(model_path: str) -> Any: From e0d44c02809cf8a0f08430a2f77c9fea892b7859 Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:32:22 +0000 Subject: [PATCH 2/6] Size the TRT-LLM capture page table with a ceiling division MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `init_capture_graph` sized the captured page table at `max_seq_len // page_size` while `_get_page_offsets` built replay tables with a ceiling division. Whenever `max_seq_len` is not a multiple of `page_size` the two disagree by one column, and the replay copy fails on a sequence long enough to need that last page: RuntimeError: The expanded size of the tensor (132) must match the existing size (133) at non-singleton dimension 1 The floor is simply wrong, not merely inconsistent: a 8505-token sequence occupies 133 pages of 64, and the engine will serve sequences up to `max_seq_len`. Route both sites through one `_num_pages` helper so a future change cannot reintroduce the skew, and reject an over-wide replay table explicitly — slicing it to fit would silently narrow the destination and drop the pages that overflow. Reachable on any model whose maximum sequence length is not page-aligned, which for Qwen3.5 is the ordinary case: 8192 input plus a 313-token allowance. --- python/minisgl/attention/trtllm.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/python/minisgl/attention/trtllm.py b/python/minisgl/attention/trtllm.py index 49144b2..9e9cd2e 100644 --- a/python/minisgl/attention/trtllm.py +++ b/python/minisgl/attention/trtllm.py @@ -58,8 +58,18 @@ def _get_decode_cu_seqlens_q(self, padded_size: int) -> torch.Tensor: ) return self._decode_cu_seqlens_q[:needed] + def _num_pages(self, num_tokens: int) -> int: + """Page-table columns a sequence of ``num_tokens`` tokens occupies. + + Capture and replay must derive this the same way. The captured graph + bakes in the width of `capture.page_table`, while replay copies a table + built from the batch's own longest sequence; if the two rules disagree + the copy fails with an opaque shape error far from the cause. + """ + return (int(num_tokens) + int(self.page_size) - 1) // int(self.page_size) + def _get_page_offsets(self, max_seqlen_k: int) -> torch.Tensor: - page_count = (int(max_seqlen_k) + int(self.page_size) - 1) // int(self.page_size) + page_count = self._num_pages(max_seqlen_k) offsets = self._page_offsets_by_len.get(page_count) if offsets is None: offsets = ( @@ -177,7 +187,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: assert self.capture is None, "Capture already initialized." max_bs = max(bs_list) capture = TRTLLMCaptureData.create( - max_bs, max_seq_len // self.page_size, self.kvcache.device + max_bs, self._num_pages(max_seq_len), self.kvcache.device ) self.capture = capture self.capture_bs = sorted(bs_list) @@ -201,6 +211,15 @@ def prepare_for_replay(self, batch: Batch) -> None: assert self.capture is not None and bs in self.capture_bs # cu_seqlens_q is always [0, 1, 2, ..., bs] for decode (i.e. no-op) table_len = metadata.page_table.size(1) + captured_len = self.capture.page_table.size(1) + if table_len > captured_len: + # Slicing to [:captured_len] would silently narrow the destination and + # surface as a bare shape mismatch, which says nothing about pages. + raise ValueError( + f"replay needs {table_len} page-table columns but the captured " + f"buffer has {captured_len} (page_size={self.page_size}); the " + f"capture and replay page-count rules have diverged" + ) self.capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k) self.capture.seq_lens[:bs].copy_(metadata.cache_seqlens) self.capture.page_table[:bs, :table_len].copy_(metadata.page_table) From 759359fe64df676d1e59fa5ae36afbafeb1567bc Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:35:41 +0000 Subject: [PATCH 3/6] Let the launcher scripts drive Qwen3.5 alignment runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps kept the shell entrypoints from expressing what a Qwen3.5 alignment run needs. `fastafd_server.sh` had no way to reach the AFD overlap controls, so `--afd-async-pending-steps`, `--afd-force-multi-mb-graph-overlap` and `--afd-disable-overlap` were unreachable from a script even though the server accepts them. The first is the one that matters in practice: it sets how far graph launches may run ahead of the CPU wait on the previous step's device-to-host copy, and the useful depth differs by model and topology. `fastafd_vllm_alignment.sh` assumed one conda environment and one model path for both engines. Qwen3.5 needs neither: FastAFD and vLLM live in separate environments here, and a reference may need its own view of a checkpoint. It also hardcoded 600 s readiness and scoring timeouts, which a 397B cold start on two nodes does not fit — those now read `AFD_READY_TIMEOUT`, `AFD_SAMPLE_TIMEOUT`, `VLLM_READY_TIMEOUT` and `VLLM_SCORE_TIMEOUT`. `--afd-num-mb` is threaded through so one preset can run both the mb1 baseline and the mb2 overlap case. All three serve scripts resolve CUDA the same way when the toolkit comes from a CUDA 13 conda package: headers and binaries sit at the prefix while target libraries are exposed through `$CONDA_PREFIX/lib64`, so pointing `CUDA_HOME` at `targets/sbsa-linux` — which the previous fallback effectively did — makes a cold FlashInfer JIT link against a lib64 that does not exist. The new branch is guarded on that layout actually being present, so environments without it are unaffected. --- scripts/serve/fastafd_server.sh | 35 +++++++++++ scripts/serve/minisgl_server.sh | 12 +++- scripts/serve/vllm_server.sh | 16 ++++- scripts/validate/fastafd_vllm_alignment.sh | 69 +++++++++++++++++----- 4 files changed, 116 insertions(+), 16 deletions(-) diff --git a/scripts/serve/fastafd_server.sh b/scripts/serve/fastafd_server.sh index 6155ddc..bc6fd13 100755 --- a/scripts/serve/fastafd_server.sh +++ b/scripts/serve/fastafd_server.sh @@ -32,6 +32,10 @@ Options: --cuda-graph-max-bs N Default: 8 --afd-decode-graph-bs CSV Optional static decode graph buckets --afd-num-mb N Default: 1 + --afd-async-pending-steps N Default: 3 + --afd-force-multi-mb-graph-overlap + Diagnostic opt-in for cross-step MB graph overlap + --afd-disable-overlap Disable cross-step overlap --afd-device-comm-num-sms N Default: 1 --page-size N Default: 1 --max-seq-len-override N Optional @@ -64,6 +68,9 @@ CACHE_TYPE="naive" CUDA_GRAPH_MAX_BS="8" AFD_DECODE_GRAPH_BS="" AFD_NUM_MB="1" +AFD_ASYNC_PENDING_STEPS="3" +AFD_FORCE_MULTI_MB_GRAPH_OVERLAP=0 +AFD_DISABLE_OVERLAP=0 AFD_DEVICE_COMM_NUM_SMS="1" PAGE_SIZE="1" MAX_SEQ_LEN_OVERRIDE="" @@ -96,6 +103,9 @@ while [[ $# -gt 0 ]]; do --cuda-graph-max-bs) CUDA_GRAPH_MAX_BS="$2"; shift 2 ;; --afd-decode-graph-bs) AFD_DECODE_GRAPH_BS="$2"; shift 2 ;; --afd-num-mb) AFD_NUM_MB="$2"; shift 2 ;; + --afd-async-pending-steps) AFD_ASYNC_PENDING_STEPS="$2"; shift 2 ;; + --afd-force-multi-mb-graph-overlap) AFD_FORCE_MULTI_MB_GRAPH_OVERLAP=1; shift ;; + --afd-disable-overlap) AFD_DISABLE_OVERLAP=1; shift ;; --afd-device-comm-num-sms) AFD_DEVICE_COMM_NUM_SMS="$2"; shift 2 ;; --page-size) PAGE_SIZE="$2"; shift 2 ;; --max-seq-len-override) MAX_SEQ_LEN_OVERRIDE="$2"; shift 2 ;; @@ -153,6 +163,7 @@ cmd=( --cache-type "$CACHE_TYPE" --cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS" --afd-num-mb "$AFD_NUM_MB" + --afd-async-pending-steps "$AFD_ASYNC_PENDING_STEPS" --afd-device-comm-num-sms "$AFD_DEVICE_COMM_NUM_SMS" --page-size "$PAGE_SIZE" ) @@ -172,6 +183,12 @@ fi if [[ -n "$AFD_DECODE_GRAPH_BS" ]]; then cmd+=(--afd-decode-graph-bs "$AFD_DECODE_GRAPH_BS") fi +if [[ "$AFD_FORCE_MULTI_MB_GRAPH_OVERLAP" -eq 1 ]]; then + cmd+=(--afd-force-multi-mb-graph-overlap) +fi +if [[ "$AFD_DISABLE_OVERLAP" -eq 1 ]]; then + cmd+=(--afd-disable-overlap) +fi if [[ -n "$MAX_SEQ_LEN_OVERRIDE" ]]; then cmd+=(--max-seq-len-override "$MAX_SEQ_LEN_OVERRIDE") fi @@ -181,5 +198,23 @@ if [[ -n "$EXTRA_ARGS" ]]; then fi export PYTHONUNBUFFERED=1 + +if [[ -n "${CONDA_PREFIX:-}" ]]; then + CUDA_TARGET_ROOT="${CONDA_PREFIX}/targets/sbsa-linux" + if [[ -f "${CUDA_TARGET_ROOT}/include/cuda_bf16.h" ]]; then + # General MiniSGL extensions (pynccl, Triton/TVM-FFI AOT) resolve CUDA + # from the Conda prefix. The NVFP4 module retargets FlashInfer in-process. + export CUDA_HOME="${CONDA_PREFIX}" + export CUDA_PATH="${CONDA_PREFIX}" + export CUDA_NVCC_EXECUTABLE="${CUDA_NVCC_EXECUTABLE:-${CONDA_PREFIX}/bin/nvcc}" + export FLASHINFER_NVCC="${FLASHINFER_NVCC:-${CONDA_PREFIX}/bin/nvcc}" + export PATH="${CONDA_PREFIX}/nvvm/bin:${CONDA_PREFIX}/bin:${CUDA_TARGET_ROOT}/bin:${PATH}" + if [[ -n "${FLASHINFER_WORKSPACE_BASE:-}" && \ + "${FLASHINFER_WORKSPACE_BASE##*/}" != "cuda-target-sbsa-linux" ]]; then + export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE%/}/cuda-target-sbsa-linux" + fi + fi +fi + cd "$CALIB_DIR" exec "${cmd[@]}" diff --git a/scripts/serve/minisgl_server.sh b/scripts/serve/minisgl_server.sh index 63f71a4..e81fc5b 100755 --- a/scripts/serve/minisgl_server.sh +++ b/scripts/serve/minisgl_server.sh @@ -172,7 +172,17 @@ export PYTHONUNBUFFERED=1 # first-time kernel compilation. Mirror the vLLM startup defaults so local mp # runs do not depend on the caller shell pre-exporting toolchain variables. if [[ -n "${CONDA_PREFIX:-}" ]]; then - if [[ -z "${CUDA_HOME:-}" && -x "${CONDA_PREFIX}/bin/nvcc" ]]; then + CUDA_TARGET_ROOT="${CONDA_PREFIX}/targets/sbsa-linux" + if [[ -f "${CUDA_TARGET_ROOT}/include/cuda_bf16.h" ]]; then + export CUDA_HOME="${CONDA_PREFIX}" + export CUDA_PATH="${CONDA_PREFIX}" + export FLASHINFER_NVCC="${FLASHINFER_NVCC:-${CONDA_PREFIX}/bin/nvcc}" + export PATH="${CONDA_PREFIX}/nvvm/bin:${CONDA_PREFIX}/bin:${CUDA_TARGET_ROOT}/bin:${PATH}" + if [[ -n "${FLASHINFER_WORKSPACE_BASE:-}" && \ + "${FLASHINFER_WORKSPACE_BASE##*/}" != "cuda-target-sbsa-linux" ]]; then + export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE%/}/cuda-target-sbsa-linux" + fi + elif [[ -z "${CUDA_HOME:-}" && -x "${CONDA_PREFIX}/bin/nvcc" ]]; then export CUDA_HOME="${CONDA_PREFIX}" fi if [[ -z "${CUDA_PATH:-}" && -n "${CUDA_HOME:-}" ]]; then diff --git a/scripts/serve/vllm_server.sh b/scripts/serve/vllm_server.sh index f4b7cde..b448f4f 100755 --- a/scripts/serve/vllm_server.sh +++ b/scripts/serve/vllm_server.sh @@ -99,7 +99,21 @@ export PYTHONUNBUFFERED=1 # Mirror the mini-sgl clean-shell defaults so ad-hoc runs don't depend on # the caller having pre-exported the toolchain env. if [[ -n "${CONDA_PREFIX:-}" ]]; then - if [[ -z "${CUDA_HOME:-}" && -x "${CONDA_PREFIX}/bin/nvcc" ]]; then + CUDA_TARGET_ROOT="${CONDA_PREFIX}/targets/sbsa-linux" + if [[ -f "${CUDA_TARGET_ROOT}/include/cuda_bf16.h" ]]; then + # CUDA 13 Conda packages keep headers/binaries at the prefix and expose + # target libraries through $CONDA_PREFIX/lib64. Pointing CUDA_HOME at the + # sbsa target directly makes FlashInfer link against a non-existent + # targets/sbsa-linux/lib64 directory during a cold JIT build. + export CUDA_HOME="${CONDA_PREFIX}" + export CUDA_PATH="${CONDA_PREFIX}" + export FLASHINFER_NVCC="${FLASHINFER_NVCC:-${CONDA_PREFIX}/bin/nvcc}" + export PATH="${CONDA_PREFIX}/nvvm/bin:${CONDA_PREFIX}/bin:${CUDA_TARGET_ROOT}/bin:${PATH}" + if [[ -n "${FLASHINFER_WORKSPACE_BASE:-}" && \ + "${FLASHINFER_WORKSPACE_BASE##*/}" != "cuda-target-sbsa-linux" ]]; then + export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE%/}/cuda-target-sbsa-linux" + fi + elif [[ -z "${CUDA_HOME:-}" && -x "${CONDA_PREFIX}/bin/nvcc" ]]; then export CUDA_HOME="${CONDA_PREFIX}" fi if [[ -z "${CUDA_PATH:-}" && -n "${CUDA_HOME:-}" ]]; then diff --git a/scripts/validate/fastafd_vllm_alignment.sh b/scripts/validate/fastafd_vllm_alignment.sh index 82843d3..024a5f5 100755 --- a/scripts/validate/fastafd_vllm_alignment.sh +++ b/scripts/validate/fastafd_vllm_alignment.sh @@ -15,13 +15,16 @@ This runs a sequential online AFD-vs-vLLM alignment pipeline: Options: --env NAME Expected active conda env. Default: minisgl-cuda130 + --vllm-env NAME Conda env used only for vLLM. Defaults to --env. --model PATH Model path or HF repo. Required. + --vllm-model PATH Optional vLLM-only model path. Defaults to --model. --prompt-file FILE UTF-8 prompt file. Required. --prompt-repeat N Repeat prompt set in memory. Default: 1 --batch-size N AFD online batch size. Default: 8 --max-new-tokens N Sampling max new tokens. Default: 1 --cuda-graph-max-bs N AFD decode graph cap. Default: batch size --decode-graph-bs CSV Optional static AFD decode graph buckets + --afd-num-mb N AFD microbatches per step. Default: 1 --max-batched-tokens N Static eager prefill token budget / comm buffer sizing. Default: 8192 --cache-type NAME KV cache manager strategy for AFD. Default: naive --gpus CSV CUDA_VISIBLE_DEVICES for vLLM. Default: 0,1,2,3 @@ -217,13 +220,16 @@ log_step() { } ENV_NAME="minisgl-cuda130" +VLLM_ENV_NAME="" MODEL="" +VLLM_MODEL="" PROMPT_FILE="" PROMPT_REPEAT="1" AFD_BATCH_SIZE="8" MAX_NEW_TOKENS="1" AFD_CUDA_GRAPH_MAX_BS="" AFD_DECODE_GRAPH_BS="" +AFD_NUM_MB="1" AFD_MAX_BATCHED_TOKENS="8192" AFD_MAX_RUNNING_REQUESTS="0" CACHE_TYPE="naive" @@ -264,13 +270,16 @@ VLLM_PID="" while [[ $# -gt 0 ]]; do case "$1" in --env) ENV_NAME="$2"; shift 2 ;; + --vllm-env) VLLM_ENV_NAME="$2"; shift 2 ;; --model) MODEL="$2"; shift 2 ;; + --vllm-model) VLLM_MODEL="$2"; shift 2 ;; --prompt-file) PROMPT_FILE="$2"; shift 2 ;; --prompt-repeat) PROMPT_REPEAT="$2"; shift 2 ;; --batch-size) AFD_BATCH_SIZE="$2"; shift 2 ;; --max-new-tokens) MAX_NEW_TOKENS="$2"; shift 2 ;; --cuda-graph-max-bs) AFD_CUDA_GRAPH_MAX_BS="$2"; shift 2 ;; --decode-graph-bs) AFD_DECODE_GRAPH_BS="$2"; shift 2 ;; + --afd-num-mb) AFD_NUM_MB="$2"; shift 2 ;; --max-batched-tokens) AFD_MAX_BATCHED_TOKENS="$2"; shift 2 ;; --afd-max-running-requests) AFD_MAX_RUNNING_REQUESTS="$2"; shift 2 ;; --cache-type) CACHE_TYPE="$2"; shift 2 ;; @@ -308,12 +317,22 @@ done if [[ -z "$AFD_CUDA_GRAPH_MAX_BS" ]]; then AFD_CUDA_GRAPH_MAX_BS="$AFD_BATCH_SIZE" fi +if (( AFD_NUM_MB < 1 )); then + echo "--afd-num-mb must be >= 1, got: $AFD_NUM_MB" >&2 + exit 1 +fi if [[ -z "$MODEL" || -z "$PROMPT_FILE" ]]; then echo "--model and --prompt-file are required" >&2 usage >&2 exit 1 fi +if [[ -z "$VLLM_MODEL" ]]; then + VLLM_MODEL="$MODEL" +fi +if [[ -z "$VLLM_ENV_NAME" ]]; then + VLLM_ENV_NAME="$ENV_NAME" +fi if [[ ! -f "$PROMPT_FILE" ]]; then echo "Prompt file not found: $PROMPT_FILE" >&2 @@ -411,12 +430,15 @@ cd "$CALIB_DIR" log_step "Starting online AFD server" echo " model: $MODEL" +echo " vllm_model: $VLLM_MODEL" +echo " vllm_env: $VLLM_ENV_NAME" echo " prompt_file: $PROMPT_FILE" echo " prompt_repeat: $PROMPT_REPEAT" echo " afd_batch_size: $AFD_BATCH_SIZE" echo " max_new_tokens: $MAX_NEW_TOKENS" echo " afd_cuda_graph_max_bs: $AFD_CUDA_GRAPH_MAX_BS" echo " afd_decode_graph_bs: ${AFD_DECODE_GRAPH_BS:-}" +echo " afd_num_mb: $AFD_NUM_MB" echo " afd_max_batched_tokens: $AFD_MAX_BATCHED_TOKENS" echo " afd_max_running_requests: ${AFD_MAX_RUNNING_REQUESTS:-0}" echo " afd_mlp_ep_size: $AFD_MLP_EP_SIZE" @@ -480,6 +502,7 @@ afd_start_cmd+=( --cache-type "$CACHE_TYPE" --page-size "$PAGE_SIZE" --cuda-graph-max-bs "$AFD_CUDA_GRAPH_MAX_BS" + --afd-num-mb "$AFD_NUM_MB" --afd-device-comm-num-sms "$AFD_DEVICE_COMM_NUM_SMS" ) if [[ -n "$AFD_DECODE_GRAPH_BS" ]]; then @@ -495,7 +518,12 @@ fi "${afd_start_cmd[@]}" >"$AFD_LOG" 2>&1 & AFD_PID=$! -if ! wait_for_ready "http://${SERVER_HOST}:${AFD_PORT}/v1/models" "AFD" "$AFD_LOG" "$AFD_PID"; then +if ! wait_for_ready \ + "http://${SERVER_HOST}:${AFD_PORT}/v1/models" \ + "AFD" \ + "$AFD_LOG" \ + "$AFD_PID" \ + "${AFD_READY_TIMEOUT:-900}"; then status="$?" if [[ "$status" == "2" ]]; then fail_with_logs "AFD exited before becoming ready" 1 "$AFD_LOG" @@ -515,7 +543,7 @@ sample_cmd=( --top-p 1 --top-k 1 --sample-concurrency "$SAMPLE_CONCURRENCY" - --timeout 600 + --timeout "${AFD_SAMPLE_TIMEOUT:-600}" --sample-json "$SAMPLE_JSON" ) if [[ "$BATCH_SUBMIT" -eq 1 ]]; then @@ -537,19 +565,32 @@ wait "$AFD_PID" >/dev/null 2>&1 || true AFD_PID="" log_step "Starting vLLM" -"$SCRIPT_ROOT/serve/vllm_server.sh" \ - --env "$ENV_NAME" \ - --model "$MODEL" \ - --gpus "$GPUS" \ - --port "$VLLM_PORT" \ - --host "$SERVER_HOST" \ - --tp-size "$VLLM_TP_SIZE" \ - --max-model-len "$VLLM_MAX_MODEL_LEN" \ - --extra-args "$VLLM_EXTRA_ARGS" \ - >"$VLLM_LOG" 2>&1 & +vllm_start_cmd=( + "$SCRIPT_ROOT/serve/vllm_server.sh" + --env "$VLLM_ENV_NAME" + --model "$VLLM_MODEL" + --gpus "$GPUS" + --port "$VLLM_PORT" + --host "$SERVER_HOST" + --tp-size "$VLLM_TP_SIZE" + --max-model-len "$VLLM_MAX_MODEL_LEN" + --extra-args "$VLLM_EXTRA_ARGS" +) +if [[ "$VLLM_ENV_NAME" != "$ENV_NAME" ]]; then + vllm_start_cmd=( + conda run --no-capture-output -n "$VLLM_ENV_NAME" + "${vllm_start_cmd[@]}" + ) +fi +"${vllm_start_cmd[@]}" >"$VLLM_LOG" 2>&1 & VLLM_PID=$! -if ! wait_for_ready "http://${SERVER_HOST}:${VLLM_PORT}/v1/models" "vLLM" "$VLLM_LOG" "$VLLM_PID"; then +if ! wait_for_ready \ + "http://${SERVER_HOST}:${VLLM_PORT}/v1/models" \ + "vLLM" \ + "$VLLM_LOG" \ + "$VLLM_PID" \ + "${VLLM_READY_TIMEOUT:-900}"; then status="$?" if [[ "$status" == "2" ]]; then fail_with_logs "vLLM exited before becoming ready" 1 "$VLLM_LOG" @@ -562,7 +603,7 @@ if ! python "$SCRIPT_ROOT/validate/compare_minisgl_vllm.py" score \ --sample-json "$SAMPLE_JSON" \ --vllm-url "http://${SERVER_HOST}:${VLLM_PORT}" \ --prompt-logprobs "$PROMPT_LOGPROBS" \ - --timeout 600 \ + --timeout "${VLLM_SCORE_TIMEOUT:-600}" \ --output-json "$REPORT_JSON" \ >"$TMP_DIR/score.log" 2>&1; then fail_with_logs "vLLM scoring failed" 1 "$TMP_DIR/score.log" "$VLLM_LOG" From 348cc19988b71b606253f4a19d2f939c72765852 Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:38:40 +0000 Subject: [PATCH 4/6] Add Qwen3.5 FP8 correctness alignment presets Two presets pin a checkpoint revision and hand everything else to a shared driver, matching how the Qwen3-30B presets are laid out. Both compare FastAFD against a vLLM reference on the same weights: FastAFD on AG-TP4 plus EG-TP/EP4 across two nodes, vLLM at TP4 in its own conda environment, prompt logprobs at top-10 over the pinned prompt set. The reference deliberately runs with `--enforce-eager`, `--language-model-only` and allreduce-RMS fusion off. A fused or graph-captured reference turns any mismatch into a question about which pass changed the arithmetic; an eager one points at the kernel under test. The driver refuses to run without MODEL_REPO, MODEL_REVISION and MODEL_DISPLAY_NAME rather than defaulting to a checkpoint. A default there would quietly download and align a model the caller did not ask for, and the pinned revision is the part that makes a result reproducible. 397B differs from 122B only in identity and timeouts: 379 GB of weights loaded on both engines does not fit the 1800 s readiness window that suits 122B. --- prompts/qwen35_122b_alignment.txt | 8 + .../qwen35_122b_a10b_fp8_fastafd_alignment.sh | 36 ++++ .../qwen35_397b_a17b_fp8_fastafd_alignment.sh | 43 ++++ scripts/validate/qwen35_fastafd_alignment.sh | 187 ++++++++++++++++++ 4 files changed, 274 insertions(+) create mode 100644 prompts/qwen35_122b_alignment.txt create mode 100755 scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh create mode 100755 scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh create mode 100755 scripts/validate/qwen35_fastafd_alignment.sh diff --git a/prompts/qwen35_122b_alignment.txt b/prompts/qwen35_122b_alignment.txt new file mode 100644 index 0000000..c28b329 --- /dev/null +++ b/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. diff --git a/scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh b/scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh new file mode 100755 index 0000000..71fbffb --- /dev/null +++ b/scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + cat <<'EOF' +Usage: + qwen35_122b_a10b_fp8_fastafd_alignment.sh + +Run the pinned official Qwen3.5-122B-A10B-FP8 two-node FastAFD/vLLM alignment. +This preset supplies only the model identity and report name; every control is +documented in the shared driver, qwen35_fastafd_alignment.sh --help. +EOF + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" + +MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3.5-122B-A10B-FP8}" +MODEL_REVISION="${MODEL_REVISION:-a099dee70ccfcd8d5dda56aaa0b60cb8ecadabc9}" +MODEL_DISPLAY_NAME="${MODEL_DISPLAY_NAME:-Qwen3.5-122B-A10B-FP8}" + +REPORT_DIR="$CALIB_DIR/reports" +mkdir -p "$REPORT_DIR" +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" +OUTPUT_JSON="${OUTPUT_JSON:-$REPORT_DIR/qwen35_122b_a10b_fp8_afd_serve_alignment_tp8_ray_${TIMESTAMP}.json}" + +export MODEL_REPO MODEL_REVISION MODEL_DISPLAY_NAME OUTPUT_JSON + +exec "$SCRIPT_DIR/qwen35_fastafd_alignment.sh" "$@" diff --git a/scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh b/scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh new file mode 100755 index 0000000..7482a56 --- /dev/null +++ b/scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + cat <<'EOF' +Usage: + qwen35_397b_a17b_fp8_fastafd_alignment.sh + +Run the pinned official Qwen3.5-397B-A17B-FP8 two-node FastAFD/vLLM alignment. +Same driver as the 122B preset, differing in model identity, report name, and +cold-start timeouts: loading 379 GB of weights on both engines does not fit the +default 1800 s readiness window. Every control is documented in +qwen35_fastafd_alignment.sh --help. +EOF + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" + +MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3.5-397B-A17B-FP8}" +MODEL_REVISION="${MODEL_REVISION:-ea5b4f81096f3901c91dea97f81324302495781d}" +MODEL_DISPLAY_NAME="${MODEL_DISPLAY_NAME:-Qwen3.5-397B-A17B-FP8}" + +REPORT_DIR="$CALIB_DIR/reports" +mkdir -p "$REPORT_DIR" +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" +OUTPUT_JSON="${OUTPUT_JSON:-$REPORT_DIR/qwen35_397b_a17b_fp8_afd_serve_alignment_tp8_ray_${TIMESTAMP}.json}" +AFD_READY_TIMEOUT="${AFD_READY_TIMEOUT:-3600}" +AFD_SAMPLE_TIMEOUT="${AFD_SAMPLE_TIMEOUT:-1200}" +VLLM_READY_TIMEOUT="${VLLM_READY_TIMEOUT:-3600}" +VLLM_SCORE_TIMEOUT="${VLLM_SCORE_TIMEOUT:-1200}" + +export MODEL_REPO MODEL_REVISION MODEL_DISPLAY_NAME OUTPUT_JSON +export AFD_READY_TIMEOUT AFD_SAMPLE_TIMEOUT VLLM_READY_TIMEOUT VLLM_SCORE_TIMEOUT + +exec "$SCRIPT_DIR/qwen35_fastafd_alignment.sh" "$@" diff --git a/scripts/validate/qwen35_fastafd_alignment.sh b/scripts/validate/qwen35_fastafd_alignment.sh new file mode 100755 index 0000000..71a5d42 --- /dev/null +++ b/scripts/validate/qwen35_fastafd_alignment.sh @@ -0,0 +1,187 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + cat <<'EOF' +Usage: + MODEL_REPO=... MODEL_REVISION=... MODEL_DISPLAY_NAME=... \ + qwen35_fastafd_alignment.sh + +Shared driver for the Qwen3.5 two-node FastAFD/vLLM alignment. It does not name +a checkpoint; a preset supplies the model identity and any model-specific +overrides, then execs this script. See: + qwen35_122b_a10b_fp8_fastafd_alignment.sh + qwen35_397b_a17b_fp8_fastafd_alignment.sh + +Configuration is supplied through the environment. The principal controls are +MAX_TOKENS, AFD_BATCH_SIZE, AFD_NUM_MB, MINISGL_CUDA_GRAPH_MAX_BS, +AFD_DECODE_GRAPH_BS, AFD_BATCH_SUBMIT, and AFD_SERVER_EXTRA_ARGS. + +The default is the eager graph-off mb1 correctness baseline. Set +AFD_SERVER_EXTRA_ARGS to an explicit empty string to enable both AFD decode +graphs once graph-off correctness passes. +EOF + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" + +# Model identity has no default on purpose: a wrong default would silently +# download and align the wrong checkpoint, and the pinned revision is what makes +# an alignment result reproducible. +for _required in MODEL_REPO MODEL_REVISION MODEL_DISPLAY_NAME; do + if [[ -z "${!_required:-}" ]]; then + echo "$_required is required. Run a preset such as" >&2 + echo " scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh" >&2 + echo "or export MODEL_REPO, MODEL_REVISION and MODEL_DISPLAY_NAME yourself." >&2 + exit 1 + fi +done +if [[ -z "${MODEL_PATH:-}" ]]; then + MODEL_PATH="$(python - "$MODEL_REPO" "$MODEL_REVISION" <<'PY' +import sys +from huggingface_hub import snapshot_download + +print(snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2])) +PY +)" +fi + +# Both engines read the same snapshot by default. A preset overrides this when a +# checkpoint's quantization metadata asks vLLM for something FastAFD does not +# want, in which case it points the reference at its own view of the snapshot. +VLLM_MODEL_PATH="${VLLM_MODEL_PATH:-$MODEL_PATH}" + +REPORT_DIR="$CALIB_DIR/reports" +mkdir -p "$REPORT_DIR" +PROMPT_FILE="${PROMPT_FILE:-$CALIB_DIR/prompts/qwen35_122b_alignment.txt}" +ENV_NAME="${ENV_NAME:-minisgl-cuda130-m2n}" +VLLM_ENV_NAME="${VLLM_ENV_NAME:-vllm-cuda130}" +GPUS="${GPUS:-0,1,2,3}" +MINISGL_TP_SIZE="${MINISGL_TP_SIZE:-8}" +ATTN_TP_SIZE="${ATTN_TP_SIZE:-4}" +MLP_TP_SIZE="${MLP_TP_SIZE:-4}" +ATTN_DP_SIZE="${ATTN_DP_SIZE:-1}" +MLP_DP_SIZE="${MLP_DP_SIZE:-1}" +AFD_MLP_EP_SIZE="${AFD_MLP_EP_SIZE:-4}" +AFD_MOE_A2A_BACKEND="${AFD_MOE_A2A_BACKEND:-none}" +AFD_MOE_RUNNER_BACKEND="${AFD_MOE_RUNNER_BACKEND:-deep_gemm}" +VLLM_TP_SIZE="${VLLM_TP_SIZE:-4}" +SERVER_HOST="${SERVER_HOST:-127.0.0.1}" +MINISGL_PORT="${MINISGL_PORT:-}" +VLLM_PORT="${VLLM_PORT:-}" +MAX_TOKENS="${MAX_TOKENS:-16}" +PROMPT_LOGPROBS="${PROMPT_LOGPROBS:-10}" +PROMPT_REPEAT="${PROMPT_REPEAT:-1}" +SAMPLE_CONCURRENCY="${SAMPLE_CONCURRENCY:-1}" +AFD_BATCH_SUBMIT="${AFD_BATCH_SUBMIT:-1}" +MINISGL_MAX_RUNNING_REQUESTS="${MINISGL_MAX_RUNNING_REQUESTS:-}" +MINISGL_CUDA_GRAPH_MAX_BS="${MINISGL_CUDA_GRAPH_MAX_BS:-1}" +AFD_DECODE_GRAPH_BS="${AFD_DECODE_GRAPH_BS:-}" +AFD_NUM_MB="${AFD_NUM_MB:-1}" +MINISGL_MAX_SEQ_LEN="${MINISGL_MAX_SEQ_LEN:-2048}" +MINISGL_PAGE_SIZE="${MINISGL_PAGE_SIZE:-1}" +MINISGL_RAY_ADDRESS="${MINISGL_RAY_ADDRESS:-auto}" +AFD_BATCH_SIZE="${AFD_BATCH_SIZE:-1}" +AFD_MAX_BATCHED_TOKENS="${AFD_MAX_BATCHED_TOKENS:-2048}" +AFD_CACHE_TYPE="${AFD_CACHE_TYPE:-naive}" +AFD_DEVICE_COMM_NUM_SMS="${AFD_DEVICE_COMM_NUM_SMS:-4}" +VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-2048}" +# The reference runs eager and unfused so that a mismatch points at a kernel +# rather than at a fusion pass, and `--language-model-only` keeps the comparison +# to the weights FastAFD actually loads. +if [[ ! -v VLLM_EXTRA_ARGS ]]; then + VLLM_EXTRA_ARGS='--enforce-eager --language-model-only --moe-backend triton --compilation-config {"pass_config":{"fuse_allreduce_rms":false}}' +fi +if [[ ! -v AFD_SERVER_EXTRA_ARGS ]]; then + AFD_SERVER_EXTRA_ARGS='--afd-disable-attention-decode-graph --afd-disable-model-decode-graph' +fi +AFD_READY_TIMEOUT="${AFD_READY_TIMEOUT:-1800}" +AFD_SAMPLE_TIMEOUT="${AFD_SAMPLE_TIMEOUT:-600}" +VLLM_READY_TIMEOUT="${VLLM_READY_TIMEOUT:-1800}" +VLLM_SCORE_TIMEOUT="${VLLM_SCORE_TIMEOUT:-600}" +export VLLM_EXTRA_ARGS AFD_SERVER_EXTRA_ARGS AFD_READY_TIMEOUT AFD_SAMPLE_TIMEOUT +export VLLM_READY_TIMEOUT VLLM_SCORE_TIMEOUT + +if [[ ! -f "$PROMPT_FILE" ]]; then + echo "Prompt file not found: $PROMPT_FILE" >&2 + exit 1 +fi +if [[ -z "$MINISGL_MAX_RUNNING_REQUESTS" ]]; then + _prompt_count="$(grep -cve '^[[:space:]]*$' "$PROMPT_FILE")" + MINISGL_MAX_RUNNING_REQUESTS="$(( _prompt_count * PROMPT_REPEAT ))" +fi +case "$AFD_BATCH_SUBMIT" in + 0) BATCH_SUBMIT_ARG="--no-batch-submit" ;; + 1) BATCH_SUBMIT_ARG="--batch-submit" ;; + *) echo "AFD_BATCH_SUBMIT must be 0 or 1, got: $AFD_BATCH_SUBMIT" >&2; exit 1 ;; +esac + +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" +OUTPUT_JSON="${OUTPUT_JSON:-$REPORT_DIR/qwen35_afd_serve_alignment_tp8_ray_${TIMESTAMP}.json}" + +echo "Running ${MODEL_DISPLAY_NAME} mb${AFD_NUM_MB} AFD alignment" +echo " model_repo: $MODEL_REPO" +echo " model_revision: $MODEL_REVISION" +echo " model_path: $MODEL_PATH" +echo " vllm_model_path: $VLLM_MODEL_PATH" +echo " prompts: $PROMPT_FILE" +echo " max_tokens: $MAX_TOKENS" +echo " afd_topology: AG-TP${ATTN_TP_SIZE} + EG-TP/EP${MLP_TP_SIZE}" +echo " afd_moe_runner_backend: $AFD_MOE_RUNNER_BACKEND" +echo " afd_cuda_graph_max_bs: $MINISGL_CUDA_GRAPH_MAX_BS" +echo " afd_decode_graph_bs: ${AFD_DECODE_GRAPH_BS:-}" +echo " afd_num_mb: $AFD_NUM_MB" +echo " afd_batch_submit: $AFD_BATCH_SUBMIT" +echo " sample_concurrency: $SAMPLE_CONCURRENCY" +echo " afd_server_extra_args: ${AFD_SERVER_EXTRA_ARGS:-}" +echo " afd_ready_timeout: $AFD_READY_TIMEOUT" +echo " vllm_tp_size: $VLLM_TP_SIZE" +echo " vllm_env: $VLLM_ENV_NAME" +echo " vllm_ready_timeout: $VLLM_READY_TIMEOUT" +echo " output_json: $OUTPUT_JSON" + +exec "$SCRIPT_ROOT/validate/fastafd_vllm_alignment.sh" \ + --env "$ENV_NAME" \ + --vllm-env "$VLLM_ENV_NAME" \ + --model "$MODEL_PATH" \ + --vllm-model "$VLLM_MODEL_PATH" \ + --prompt-file "$PROMPT_FILE" \ + --prompt-repeat "$PROMPT_REPEAT" \ + --batch-size "$AFD_BATCH_SIZE" \ + --max-new-tokens "$MAX_TOKENS" \ + --cuda-graph-max-bs "$MINISGL_CUDA_GRAPH_MAX_BS" \ + --decode-graph-bs "$AFD_DECODE_GRAPH_BS" \ + --afd-num-mb "$AFD_NUM_MB" \ + --max-batched-tokens "$AFD_MAX_BATCHED_TOKENS" \ + --cache-type "$AFD_CACHE_TYPE" \ + --afd-max-running-requests "$MINISGL_MAX_RUNNING_REQUESTS" \ + --gpus "$GPUS" \ + --afd-attn-dp-size "$ATTN_DP_SIZE" \ + --afd-mlp-dp-size "$MLP_DP_SIZE" \ + --attn-tp-size "$ATTN_TP_SIZE" \ + --mlp-tp-size "$MLP_TP_SIZE" \ + --afd-mlp-ep-size "$AFD_MLP_EP_SIZE" \ + --afd-moe-a2a-backend "$AFD_MOE_A2A_BACKEND" \ + --afd-moe-runner-backend "$AFD_MOE_RUNNER_BACKEND" \ + --vllm-tp-size "$VLLM_TP_SIZE" \ + --ray-address "$MINISGL_RAY_ADDRESS" \ + --host "$SERVER_HOST" \ + --afd-port "$MINISGL_PORT" \ + --vllm-port "$VLLM_PORT" \ + --prompt-logprobs "$PROMPT_LOGPROBS" \ + --sample-concurrency "$SAMPLE_CONCURRENCY" \ + "$BATCH_SUBMIT_ARG" \ + --max-seq-len-override "$MINISGL_MAX_SEQ_LEN" \ + --page-size "$MINISGL_PAGE_SIZE" \ + --vllm-max-model-len "$VLLM_MAX_MODEL_LEN" \ + --output-json "$OUTPUT_JSON" \ + --afd-device-comm-num-sms "$AFD_DEVICE_COMM_NUM_SMS" \ + --keep-artifacts From 2fb1d58e444b3607c35db84ea6b9e94391d701b4 Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:44:18 +0000 Subject: [PATCH 5/6] Add Qwen3.5 two-node AFD throughput presets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One preset per model over a shared no-TP driver: attention DP4/TP1 and expert DP4/TP1/EP4 across two nodes, MegaMoE M:N transport, 8k prompts replayed from the pinned pool with Nsight capturing a decode window. The driver refuses three things rather than guessing them. Tensor parallelism is rejected on both roles, so a retained decode window contains no TP all-reduce kernels and a transport measurement is not contaminated by one. `AFD_NUM_MB` has no default, because 1 is the serialized baseline and 2 is the overlap case and picking either silently turns a throughput comparison into a coin toss. And the Nsight stop step is checked against the minimum reachable total step: overshoot it and the worker sessions never finalize, leaving only coordinator traces — which looks like a profiling bug rather than a window that was too wide. Batch is named per attention GPU, as the 235B presets do: 256 for 122B and 320 for 397B. 397B is 60 layers to 122B's 48 and routes 512 experts at top-k 10, so the same eight GPUs hold fewer requests. Prefill duration is bracketed rather than point-estimated. Dividing prompt tokens by the token cap assumes no decode contention and is a lower bound; requests that finish prefill early consume a decode slot on every later iteration, shrinking the per-step prefill budget toward (cap - active), which gives the upper bound. The capture window spans the bracket and the analyser isolates the pure-decode tail by kernel timestamp. Verified by comparing each preset's derived configuration against the driver the measurements were taken with: all 22 echoed values match at both batch sizes. --- ...0b_fp8_8k_b256_2node_mb2_nsys_alignment.sh | 33 +++ ...7b_fp8_8k_b320_2node_mb2_nsys_alignment.sh | 35 +++ .../run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh | 269 ++++++++++++++++++ 3 files changed, 337 insertions(+) create mode 100755 scripts/experiments/afd/qwen35/run_afd_qwen35_122b_a10b_fp8_8k_b256_2node_mb2_nsys_alignment.sh create mode 100755 scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh create mode 100755 scripts/experiments/afd/qwen35/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh diff --git a/scripts/experiments/afd/qwen35/run_afd_qwen35_122b_a10b_fp8_8k_b256_2node_mb2_nsys_alignment.sh b/scripts/experiments/afd/qwen35/run_afd_qwen35_122b_a10b_fp8_8k_b256_2node_mb2_nsys_alignment.sh new file mode 100755 index 0000000..77d1ac6 --- /dev/null +++ b/scripts/experiments/afd/qwen35/run_afd_qwen35_122b_a10b_fp8_8k_b256_2node_mb2_nsys_alignment.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" + +# Qwen3.5-122B-A10B-FP8 AFD AG/EG two-node 8k decode replay workload. +# +# A thin wrapper around the shared no-TP driver, which owns the topology, the +# MegaMoE M:N settings, the prompt cycling, the Nsight window and cleanup. This +# file only names the model and the workload shape. +# +# 256 concurrently decoding requests per attention GPU across four attention DP +# replicas. With mb=2 that replays a 128-request per-microbatch graph. Results +# were measured against revision a099dee70ccfcd8d5dda56aaa0b60cb8ecadabc9. +export MODEL_PATH="${MODEL_PATH:-Qwen/Qwen3.5-122B-A10B-FP8}" +export PROMPT_LEN="${PROMPT_LEN:-8192}" +export PER_ATTN_GPU_BSZ="${PER_ATTN_GPU_BSZ:-256}" +export AFD_NUM_MB="${AFD_NUM_MB:-2}" + +export ATTN_DP_SIZE="${ATTN_DP_SIZE:-4}" +export AFD_ACTIVE_GLOBAL="${AFD_ACTIVE_GLOBAL:-$((PER_ATTN_GPU_BSZ * ATTN_DP_SIZE))}" + +export RUN_DIR="${RUN_DIR:-$CALIB_DIR/reports/afd_qwen35_122b_a10b_fp8_8k_b${PER_ATTN_GPU_BSZ}_peragpu_2node_mb${AFD_NUM_MB}_${TIMESTAMP}}" + +exec bash "$SCRIPT_DIR/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh" diff --git a/scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh b/scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh new file mode 100755 index 0000000..8cc92a0 --- /dev/null +++ b/scripts/experiments/afd/qwen35/run_afd_qwen35_397b_a17b_fp8_8k_b320_2node_mb2_nsys_alignment.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" + +# Qwen3.5-397B-A17B-FP8 AFD AG/EG two-node 8k decode replay workload. +# +# A thin wrapper around the shared no-TP driver, which owns the topology, the +# MegaMoE M:N settings, the prompt cycling, the Nsight window and cleanup. This +# file only names the model and the workload shape. +# +# 320 concurrently decoding requests per attention GPU across four attention DP +# replicas. With mb=2 that replays a 160-request per-microbatch graph. 397B is +# 60 layers against 122B's 48 and routes 512 experts at top-k 10, so the same +# GPUs hold fewer requests than the 122B preset's 256. Results were measured +# against revision ea5b4f81096f3901c91dea97f81324302495781d. +export MODEL_PATH="${MODEL_PATH:-Qwen/Qwen3.5-397B-A17B-FP8}" +export PROMPT_LEN="${PROMPT_LEN:-8192}" +export PER_ATTN_GPU_BSZ="${PER_ATTN_GPU_BSZ:-320}" +export AFD_NUM_MB="${AFD_NUM_MB:-2}" + +export ATTN_DP_SIZE="${ATTN_DP_SIZE:-4}" +export AFD_ACTIVE_GLOBAL="${AFD_ACTIVE_GLOBAL:-$((PER_ATTN_GPU_BSZ * ATTN_DP_SIZE))}" + +export RUN_DIR="${RUN_DIR:-$CALIB_DIR/reports/afd_qwen35_397b_a17b_fp8_8k_b${PER_ATTN_GPU_BSZ}_peragpu_2node_mb${AFD_NUM_MB}_${TIMESTAMP}}" + +exec bash "$SCRIPT_DIR/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh" diff --git a/scripts/experiments/afd/qwen35/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh b/scripts/experiments/afd/qwen35/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh new file mode 100755 index 0000000..87966bb --- /dev/null +++ b/scripts/experiments/afd/qwen35/run_qwen35_fp8_afd_dp4_ep4_notp_2node.sh @@ -0,0 +1,269 @@ +#!/usr/bin/env bash +# Shared Qwen3.5 FP8 two-node AFD throughput driver: AG DP4/TP1 + EG DP4/TP1/EP4, +# MegaMoE M:N transport. A preset supplies MODEL_PATH, the batch and the +# microbatch count, then execs this script. +# +# Four properties this wrapper enforces rather than assumes: +# 1. tensor parallelism is forbidden on both roles, so a retained decode window +# contains zero TP all-reduce kernels; +# 2. AFD_NUM_MB must be explicit -- 1 is the serialized baseline and 2 is the +# overlap case, and silently defaulting either way makes a throughput +# comparison meaningless; +# 3. the workload is the pinned prompt pool, deterministically cycled by request +# index, so AFD and a vLLM reference decode byte-identical prompts; +# 4. batches are named globally (AFD_ACTIVE_GLOBAL) and the per-rank and +# per-lane shapes are derived and echoed rather than implied. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +_common_dir="$SCRIPT_DIR" +while [[ "$_common_dir" != "/" && ! -f "$_common_dir/scripts/lib/common.sh" ]]; do + _common_dir="$(dirname "$_common_dir")" +done +# shellcheck source=/dev/null +source "$_common_dir/scripts/lib/common.sh" +fastafd_init_paths "$SCRIPT_DIR" + +RUNNER="$SCRIPT_ROOT/experiments/afd/qwen3_30b/run_afd_qwen3_30b_a3b_fp8_3node_mb2_nsys_alignment.sh" +TIMESTAMP="$(date -u +%Y%m%d_%H%M%S)" + +if [[ -z "${MODEL_PATH:-}" ]]; then + echo "MODEL_PATH must name a pinned Qwen3.5 FP8 snapshot." >&2 + exit 1 +fi + +export ENV_NAME="${ENV_NAME:-minisgl-cuda130-m2n}" +export PROMPT_LEN="${PROMPT_LEN:-8192}" +# A longer decode plateau makes the retained window robust to prefill-duration +# estimation error; decode steps are identical regardless of how many run. +export MAX_TOKENS="${MAX_TOKENS:-48}" +export AFD_MAX_NEW_TOKENS="${AFD_MAX_NEW_TOKENS:-$MAX_TOKENS}" +export WARMUP_REPLAYS="${WARMUP_REPLAYS:-5}" + +# ---- no tensor parallelism on either role ----------------------------------- +export ATTN_DP_SIZE="${ATTN_DP_SIZE:-4}" +export ATTN_TP_SIZE="${ATTN_TP_SIZE:-1}" +export MLP_DP_SIZE="${MLP_DP_SIZE:-4}" +export MLP_TP_SIZE="${MLP_TP_SIZE:-1}" +export AFD_MLP_EP_SIZE="${AFD_MLP_EP_SIZE:-4}" +if ((ATTN_TP_SIZE != 1 || MLP_TP_SIZE != 1)); then + echo "This wrapper forbids tensor parallelism; got ATTN_TP_SIZE=$ATTN_TP_SIZE MLP_TP_SIZE=$MLP_TP_SIZE." >&2 + exit 1 +fi +_ATTN_WORKERS=$((ATTN_DP_SIZE * ATTN_TP_SIZE)) +_MLP_WORKERS=$((MLP_DP_SIZE * MLP_TP_SIZE)) +if ((_ATTN_WORKERS != 4 || _MLP_WORKERS != 4)); then + echo "Two-node wrapper requires four AG and four EG workers; got AG=$_ATTN_WORKERS EG=$_MLP_WORKERS." >&2 + exit 1 +fi + +# ---- microbatch lanes must be explicit -------------------------------------- +if [[ -z "${AFD_NUM_MB+x}" ]]; then + echo "AFD_NUM_MB must be set explicitly (1 for the serialized baseline, 2 for the overlap case)." >&2 + exit 1 +fi +export AFD_NUM_MB +if ((AFD_NUM_MB < 1)); then + echo "AFD_NUM_MB must be >= 1; got $AFD_NUM_MB." >&2 + exit 1 +fi + +# ---- global batch namespace -------------------------------------------------- +# AFD_ACTIVE_GLOBAL is the target number of concurrently decoding requests across +# all attention DP replicas. Per-rank and per-lane shapes are derived from it. +export AFD_ACTIVE_GLOBAL="${AFD_ACTIVE_GLOBAL:-1024}" +_GRAPH_LANES=$((ATTN_DP_SIZE * AFD_NUM_MB)) +if ((AFD_ACTIVE_GLOBAL % _GRAPH_LANES != 0)); then + echo "AFD_ACTIVE_GLOBAL=$AFD_ACTIVE_GLOBAL must be divisible by ATTN_DP_SIZE*AFD_NUM_MB=$_GRAPH_LANES." >&2 + exit 1 +fi +_B_PER_AG_RANK=$((AFD_ACTIVE_GLOBAL / ATTN_DP_SIZE)) +_B_PER_LANE=$((AFD_ACTIVE_GLOBAL / _GRAPH_LANES)) + +export AFD_MAX_BATCHED_TOKENS="${AFD_MAX_BATCHED_TOKENS:-2048}" +# Chunked prefill leaves one request per attention DP replica behind the first +# decode wave, so submit one spare per replica to reach the target active wave. +if ((PROMPT_LEN > AFD_MAX_BATCHED_TOKENS)); then + _SPARE_PROMPTS="$ATTN_DP_SIZE" +else + _SPARE_PROMPTS=0 +fi +export NUM_PROMPTS="${NUM_PROMPTS:-$((AFD_ACTIVE_GLOBAL + _SPARE_PROMPTS))}" + +# ---- overlap policy ---------------------------------------------------------- +if [[ -z "${AFD_DISABLE_OVERLAP+x}" ]]; then + if ((AFD_NUM_MB == 1)); then AFD_DISABLE_OVERLAP=1; else AFD_DISABLE_OVERLAP=0; fi +fi +export AFD_DISABLE_OVERLAP +if [[ -z "${AFD_FORCE_MULTI_MB_GRAPH_OVERLAP+x}" ]]; then + # Multi-MB decode graphs are auto-serialized by the collective re-entry guard; + # cross-step overlap is a diagnostic opt-in and is always echoed below. + if ((AFD_NUM_MB > 1)); then + AFD_FORCE_MULTI_MB_GRAPH_OVERLAP=1 + else + AFD_FORCE_MULTI_MB_GRAPH_OVERLAP=0 + fi +fi +export AFD_FORCE_MULTI_MB_GRAPH_OVERLAP +export AFD_ASYNC_PENDING_STEPS="${AFD_ASYNC_PENDING_STEPS:-5}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-32}" + +export AFD_MOE_A2A_BACKEND="${AFD_MOE_A2A_BACKEND:-none}" +export AFD_MOE_RUNNER_BACKEND="${AFD_MOE_RUNNER_BACKEND:-deep_gemm}" +export MINISGL_AFD_MOE_BACKEND="${MINISGL_AFD_MOE_BACKEND:-megamoe_m2n}" + +export AFD_BATCH_SIZE="${AFD_BATCH_SIZE:-$NUM_PROMPTS}" +export SAMPLE_CONCURRENCY="${SAMPLE_CONCURRENCY:-$NUM_PROMPTS}" +export MINISGL_MAX_RUNNING_REQUESTS="${MINISGL_MAX_RUNNING_REQUESTS:-$NUM_PROMPTS}" +# MINISGL_MAX_SEQ_LEN is deliberately NOT set here: it depends on the generated +# token count, which is derived from the prefill bracket further down. Sizing it +# as PROMPT_LEN+64 caps decode at ~64 tokens, so requests retire at staggered +# points, the graph replays a fixed batch full of PADDING, and attention looks +# ~38x cheap while the real active batch has drained. See the export below. +export MINISGL_PAGE_SIZE="${MINISGL_PAGE_SIZE:-64}" + +_GRAPH_MB_BS="$_B_PER_LANE" +if [[ -z "${AFD_DECODE_GRAPH_BS:-}" ]]; then + if ((_GRAPH_MB_BS > 1)); then + # Prime the shared M:N runtime with the smallest graph before capturing the + # target graph; a cold first dispatch at the target shape can fail readiness. + export AFD_DECODE_GRAPH_BS="1,$_GRAPH_MB_BS" + else + export AFD_DECODE_GRAPH_BS="1" + fi +fi +export MINISGL_CUDA_GRAPH_MAX_BS="${MINISGL_CUDA_GRAPH_MAX_BS:-$_GRAPH_MB_BS}" +export AFD_MEMORY_RATIO="${AFD_MEMORY_RATIO:-0.78}" +export AFD_SERVER_EXTRA_ARGS="${AFD_SERVER_EXTRA_ARGS:---attention-backend trtllm --memory-ratio $AFD_MEMORY_RATIO}" + +# ---- pinned workload, deterministically cycled ------------------------------ +# The vLLM reference harness cycles the pool by request_id % pool_rows. +# Materialize the same sequence so both frameworks decode identical prompts. +_PROMPT_POOL="${PROMPT_POOL:-$CALIB_DIR/prompts/prompts_512x${PROMPT_LEN}_seed20260527.txt}" +if [[ ! -f "$_PROMPT_POOL" ]]; then + echo "Prompt pool not found: $_PROMPT_POOL" >&2 + exit 1 +fi +_PROMPT_CACHE="$CALIB_DIR/reports/prompts_cycled_${NUM_PROMPTS}x${PROMPT_LEN}.txt" +if [[ -z "${PROMPT_FILE:-}" ]]; then + if [[ ! -f "$_PROMPT_CACHE" ]]; then + mkdir -p "$(dirname "$_PROMPT_CACHE")" + awk -v want="$NUM_PROMPTS" ' + { pool[n++] = $0 } + END { + if (n == 0) { print "empty prompt pool" > "/dev/stderr"; exit 1 } + for (i = 0; i < want; i++) print pool[i % n] + }' "$_PROMPT_POOL" > "$_PROMPT_CACHE" + fi + export PROMPT_FILE="$_PROMPT_CACHE" +fi +export PROMPT_KIND="${PROMPT_KIND:-existing}" +_PROMPT_POOL_SHA="$(sha256sum "$_PROMPT_POOL" | awk '{print $1}')" + +# ---- Nsight window ----------------------------------------------------------- +_PREFILL_CHUNK_TOKENS="$AFD_MAX_BATCHED_TOKENS" +if ((_PREFILL_CHUNK_TOKENS > PROMPT_LEN)); then + _PREFILL_CHUNK_TOKENS="$PROMPT_LEN" +fi +_PREFILL_CHUNKS_PER_REQ=$(((PROMPT_LEN + _PREFILL_CHUNK_TOKENS - 1) / _PREFILL_CHUNK_TOKENS)) +_PREFILL_CHUNKS_PER_STEP=$((AFD_MAX_BATCHED_TOKENS / _PREFILL_CHUNK_TOKENS)) +if ((_PREFILL_CHUNKS_PER_STEP < 1)); then + _PREFILL_CHUNKS_PER_STEP=1 +fi +_PREFILL_PROMPTS_PER_DP=$(((NUM_PROMPTS + ATTN_DP_SIZE - 1) / ATTN_DP_SIZE)) +# Prefill duration is bracketed, not point-estimated. Dividing prompt tokens by +# the token cap assumes no decode contention and is a LOWER bound; requests that +# finish prefill immediately consume a decode-token slot on every later +# iteration, so the per-step prefill budget shrinks toward (cap - active), which +# gives an UPPER bound. The Nsight window spans the bracket and the analyzer +# isolates the pure-decode tail by kernel timestamp. +_PREFILL_STEPS_LOW=$(((_PREFILL_PROMPTS_PER_DP * _PREFILL_CHUNKS_PER_REQ + _PREFILL_CHUNKS_PER_STEP - 1) / _PREFILL_CHUNKS_PER_STEP)) +_PREFILL_TOKENS_PER_DP=$((_PREFILL_PROMPTS_PER_DP * PROMPT_LEN)) +_PREFILL_BUDGET=$((AFD_MAX_BATCHED_TOKENS - _B_PER_AG_RANK)) +if ((_PREFILL_BUDGET < 1)); then + _PREFILL_BUDGET=1 +fi +_PREFILL_STEPS_HIGH=$(((_PREFILL_TOKENS_PER_DP + _PREFILL_BUDGET - 1) / _PREFILL_BUDGET)) +if ((_PREFILL_STEPS_HIGH < _PREFILL_STEPS_LOW)); then + _PREFILL_STEPS_HIGH="$_PREFILL_STEPS_LOW" +fi +# Scheduler-step model. The harness runs a warmup pass and then the formal pass, +# each of which performs a full prefill: +# formal decode starts at : 2*prefill_real + WARMUP_REPLAYS +# total steps : 2*prefill_real + WARMUP_REPLAYS + MAX_TOKENS - 1 +# with prefill_real in [LOW, HIGH]. The profiler must START at or after the LATEST +# possible decode start, and STOP at or before the EARLIEST possible total, or the +# worker sessions never finalize and only coordinator traces are written. +_DECODE_START_LATEST=$((2 * _PREFILL_STEPS_HIGH + WARMUP_REPLAYS)) +_NSYS_DECODE_STEPS="${NSYS_DECODE_STEPS:-40}" +_NSYS_START_DEFAULT="$_DECODE_START_LATEST" +if ((_NSYS_START_DEFAULT < 1)); then + _NSYS_START_DEFAULT=1 +fi +_NSYS_STOP_DEFAULT=$((_NSYS_START_DEFAULT + _NSYS_DECODE_STEPS)) +# Generate enough tokens that the stop step is reachable even in the fastest +# (LOW-bound) prefill case. Decode steps are identical regardless of how many run. +_MIN_TOKENS_FOR_STOP=$((_NSYS_STOP_DEFAULT - 2 * _PREFILL_STEPS_LOW - WARMUP_REPLAYS + 1)) +_TOKEN_MARGIN="${TOKEN_MARGIN:-32}" +_REQUIRED_TOKENS=$((_MIN_TOKENS_FOR_STOP + _TOKEN_MARGIN)) +if ((_REQUIRED_TOKENS > MAX_TOKENS)); then + MAX_TOKENS="$_REQUIRED_TOKENS" + export MAX_TOKENS + export AFD_MAX_NEW_TOKENS="$MAX_TOKENS" +fi +_DECODE_REPLAYS=$((MAX_TOKENS > 0 ? MAX_TOKENS - 1 : 0)) +_TOTAL_STEPS_MIN=$((2 * _PREFILL_STEPS_LOW + WARMUP_REPLAYS + _DECODE_REPLAYS)) + +# Every request must be able to generate all MAX_TOKENS without hitting the +# sequence ceiling, otherwise requests retire at staggered points and the decode +# graph replays a mostly-padded batch. PROMPT_TOKENS_MAX is the tokenizer-measured +# ceiling of the pinned pool (8090-8187 for the 8K lane), not the nominal name. +export PROMPT_TOKENS_MAX="${PROMPT_TOKENS_MAX:-$PROMPT_LEN}" +export MINISGL_MAX_SEQ_LEN="${MINISGL_MAX_SEQ_LEN:-$((PROMPT_TOKENS_MAX + MAX_TOKENS + 64))}" +export NSYS_START_STEP="${NSYS_START_STEP:-$_NSYS_START_DEFAULT}" +export NSYS_STOP_STEP="${NSYS_STOP_STEP:-$_NSYS_STOP_DEFAULT}" +if ((NSYS_STOP_STEP > _TOTAL_STEPS_MIN)); then + echo "Nsight stop step $NSYS_STOP_STEP exceeds the minimum reachable total step $_TOTAL_STEPS_MIN; worker traces would never finalize." >&2 + exit 1 +fi + +export NSYS="${NSYS:-1}" +export NSYS_CUDA_GRAPH_TRACE="${NSYS_CUDA_GRAPH_TRACE:-node}" +export RUN_VLLM_ALIGNMENT="${RUN_VLLM_ALIGNMENT:-0}" +export CLEAN_FIRST="${CLEAN_FIRST:-1}" +export CLEAN_AFTER="${CLEAN_AFTER:-1}" +export MINISGL_AFD_SHUTDOWN_TIMEOUT_S="${MINISGL_AFD_SHUTDOWN_TIMEOUT_S:-30}" +export AFD_STOP_TIMEOUT_S="${AFD_STOP_TIMEOUT_S:-90}" + +export RUN_DIR="${RUN_DIR:-$CALIB_DIR/reports/qwen35_fp8_${PROMPT_LEN}_B${AFD_ACTIVE_GLOBAL}_dp4ep4_notp_mb${AFD_NUM_MB}_${TIMESTAMP}}" + +echo "Qwen3.5 FP8 two-node no-TP replay workload" +echo " model: $MODEL_PATH" +echo " topology: AG-DP${ATTN_DP_SIZE}/TP${ATTN_TP_SIZE} + EG-DP${MLP_DP_SIZE}/TP${MLP_TP_SIZE}/EP${AFD_MLP_EP_SIZE} (tensor_parallel=none)" +echo " prompt_len: $PROMPT_LEN" +echo " prompt_pool: $_PROMPT_POOL" +echo " prompt_pool_sha256: $_PROMPT_POOL_SHA" +echo " prompt_file: $PROMPT_FILE (kind=$PROMPT_KIND, cycled)" +echo " submitted_prompts: $NUM_PROMPTS" +echo " active_global: $AFD_ACTIVE_GLOBAL" +echo " active_per_ag_rank: $_B_PER_AG_RANK" +echo " active_per_lane: $_B_PER_LANE" +echo " microbatches: $AFD_NUM_MB (explicit)" +echo " overlap_policy: disable=$AFD_DISABLE_OVERLAP force_multi_mb=$AFD_FORCE_MULTI_MB_GRAPH_OVERLAP pending=$AFD_ASYNC_PENDING_STEPS cuda_connections=$CUDA_DEVICE_MAX_CONNECTIONS" +echo " afd_moe_backend: $MINISGL_AFD_MOE_BACKEND" +echo " target_graph_microbatch: $_GRAPH_MB_BS" +echo " graph_warmup_batches: $AFD_DECODE_GRAPH_BS" +echo " server_extra_args: $AFD_SERVER_EXTRA_ARGS" +echo " prefill_step_bracket: $_PREFILL_STEPS_LOW..$_PREFILL_STEPS_HIGH" +echo " generated_tokens: $MAX_TOKENS (auto-raised so the stop step is reachable)" +echo " min_reachable_total_step: $_TOTAL_STEPS_MIN" +echo " prompt_tokens_max: $PROMPT_TOKENS_MAX" +echo " max_seq_len: $MINISGL_MAX_SEQ_LEN (prompt_tokens_max + generated + 64)" +echo " nsys_step_window: $NSYS_START_STEP..$NSYS_STOP_STEP" +echo " run_dir: $RUN_DIR" + +if [[ "${QWEN35_DRY_RUN:-0}" == "1" ]]; then + exit 0 +fi + +exec bash "$RUNNER" From 1a29bc7be12a7dd35f57b743ab3a7c7e6ba4d41b Mon Sep 17 00:00:00 2001 From: yuxuandexter Date: Wed, 5 Aug 2026 07:48:14 +0000 Subject: [PATCH 6/6] Document running Qwen3.5 122B and 397B MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds Qwen3.5 to Supported Models, a correctness-alignment block naming both presets, and a row per model in the large-scale experiment table. The topology note is the part a reader needs before running anything. Every other preset in that table sizes itself from the cluster through `AFD_TOTAL_NODES`; the Qwen3.5 ones do not, and refuse to start on anything but four attention and four expert workers. Without saying so, the obvious move — copying the 235B command line and its `AFD_TOTAL_NODES=4` — fails with a worker-count error that reads like a cluster problem. Batch is stated per attention GPU, matching the existing rows: 256 for 122B and 320 for 397B. No throughput figures, since none have been published. --- README.md | 32 ++++++++++++++++++++++++++++++++ scripts/README.md | 6 ++++++ 2 files changed, 38 insertions(+) diff --git a/README.md b/README.md index f4a403e..6ce4ce6 100644 --- a/README.md +++ b/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`), diff --git a/scripts/README.md b/scripts/README.md index 8a8a57d..cd8c1b6 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -27,6 +27,11 @@ Published Qwen3-235B and MiniMax M2.5 result presets live under `scripts/experiments/afd/qwen3_235b/` and `scripts/experiments/afd/minimax_m25/`; they dispatch through `scripts/run_afd_qwen3_30b_a3b_fp8_3node_mb2_nsys_alignment.sh`. +Qwen3.5 presets live under `scripts/experiments/afd/qwen35/` and dispatch through +a shared no-TP driver in the same directory. Their correctness counterparts are +`scripts/validate/qwen35_122b_a10b_fp8_fastafd_alignment.sh` and +`scripts/validate/qwen35_397b_a17b_fp8_fastafd_alignment.sh`. + ## Layout - `quickstart/`: mini-sgl sampling workflow. @@ -35,6 +40,7 @@ they dispatch through `scripts/run_afd_qwen3_30b_a3b_fp8_3node_mb2_nsys_alignmen - `experiments/afd/qwen3_30b/`: the FastAFD launcher the presets dispatch to. - `experiments/afd/qwen3_235b/`: Qwen3-235B published-result presets. - `experiments/afd/minimax_m25/`: MiniMax M2.5 published-result presets. +- `experiments/afd/qwen35/`: Qwen3.5 two-node presets and their shared driver. - `experiments/vllm/`: the sharded vLLM alignment scorer. - `data_gen/`: prompt generation helper. - `lib/`: shared shell helpers.