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/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/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/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) 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: 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. 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" 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" 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