"""DSV4 Router — token-to-expert assignment. Two routing modes that share an output shape: - 'dense': sqrt(softplus(X @ W_gate)) + per-expert bias, top-k selection. Used by MoE layers 3+ (the bulk of the network). - 'hash': deterministic per-token-ID lookup, uniform weights. Used by the first 3 MoE layers per DSV4 §2.1. Both modes produce (topk_weights, topk_ids) suitable for direct consumption by Nvfp4MoE.run(). CUDA-graph-compatible: pre-allocated buffers, no CPU-GPU syncs. Selection between modes is by layer_idx at construction time — the kernel path is fixed once the Router is built so the dispatch is constant-folded by torch.compile. """ from __future__ import annotations from typing import Optional, Literal import torch from dsv4.ops.router import ( register_router, dense_router_op, hash_router_op, ) RouterMode = Literal["dense", "hash"] class Router: """DSV4 expert router. Per the DeepSeek-V4 paper (§2.1): - Affinity activation is sqrt(softplus(·)), replacing V3's sigmoid(·). - Auxiliary-loss-free strategy: a learned per-expert bias (loaded from checkpoint, frozen at inference) is added to the activation for SELECTION only. The actual gating weight applied to expert outputs uses the UNBIASED activation. - First 3 MoE layers use Hash routing (Roller et al. 2021): a precomputed [vocab_size, k] LUT mapping token IDs to expert IDs. No gate GEMM is performed. - Sequence-wise balance loss is training-only; not applied here. Parameters ---------- hidden_size : int Model hidden dimension. Must match W_gate's K dimension. num_experts : int Total routed experts (Flash: 256, Pro: 384). Shared experts are handled separately by Nvfp4SharedExpert. top_k : int Experts activated per token. DSV4 uses 6. routed_scaling_factor : float Post-renormalization scale on gating weights. DSV3 used 2.5; verify against the V4 checkpoint config — may be per-layer. mode : {'dense', 'hash'} Routing strategy. Decided at construction; cannot change at runtime. vocab_size : int, optional Required when mode='hash'. The LUT is [vocab_size, top_k] int32. max_num_tokens : int Upper bound on N for pre-allocated buffer sizing. device : str CUDA device. """ def __init__( self, hidden_size: int, num_experts: int, top_k: int = 6, routed_scaling_factor: float = 2.5, *, mode: RouterMode, vocab_size: Optional[int] = None, max_num_tokens: int = 8192, device: str = "cuda", ): if mode == "hash" and vocab_size is None: raise ValueError("vocab_size is required when mode='hash'") if mode not in ("dense", "hash"): raise ValueError(f"unknown router mode: {mode!r}") self.hidden_size = hidden_size self.num_experts = num_experts self.top_k = top_k self.routed_scaling_factor = routed_scaling_factor self.mode = mode self.vocab_size = vocab_size self.max_num_tokens = max_num_tokens self.device = device # ---- Parameters (filled by load_weights / finalize_weights) ---- # Dense mode — fused NVFP4 kernel (single-kernel, preferred): # gate_weight: raw NVFP4 gate weight tensor [K_packed, E_packed] uint8 # gate_weight_scale: weight scale [K_sf, E_sf] FP8 E4M3 # gate_ws2: weight_scale_2 (global scale base) # gate_input_scale: input_scale (activation global scale base) # Dense mode — 2-kernel NVFP4 path (fallback): # gate_lin: Nvfp4Linear for the gate projection # Dense mode — BF16 fallback: # W_gate: BF16 weight for cuBLAS when NVFP4 scales not available # Hash mode: # hash_lut: [vocab_size, top_k] int32 — precomputed expert IDs. self.gate_weight = None # Raw NVFP4 weight for fused kernel self.gate_weight_scale = None # FP8 E4M3 scale for fused kernel self.gate_ws2 = None # weight_scale_2 for fused kernel self.gate_input_scale = None # input_scale for fused kernel self.gate_lin = None # Nvfp4Linear for 2-kernel NVFP4 path self.W_gate: Optional[torch.Tensor] = None # BF16 fallback self.e_bias: Optional[torch.Tensor] = None self.hash_lut: Optional[torch.Tensor] = None # ---- Pre-allocated output buffers (cudagraph-safe) ---- self._topk_weights_buf: Optional[torch.Tensor] = None self._topk_ids_buf: Optional[torch.Tensor] = None # Runner ID assigned on first call (see custom_op pattern). self._runner_id: Optional[int] = None # ------------------------------------------------------------------ # Weight loading # ------------------------------------------------------------------ def load_weights( self, W_gate: Optional[torch.Tensor] = None, e_bias: Optional[torch.Tensor] = None, hash_lut: Optional[torch.Tensor] = None, ) -> None: """Populate router parameters from a checkpoint shard. Dense mode expects (W_gate, e_bias). Hash mode expects (hash_lut). Mismatches with self.mode raise immediately — these errors are nearly always loader bugs and silent acceptance would mask them. """ if self.mode == "dense": if e_bias is None: raise ValueError("dense router needs e_bias") assert e_bias.shape == (self.num_experts,), \ f"e_bias shape {tuple(e_bias.shape)} != ({self.num_experts},)" self.e_bias = e_bias.to(device=self.device, dtype=torch.float32) if W_gate is not None: self.W_gate = W_gate.to(device=self.device, dtype=torch.bfloat16) # gate_lin is set separately via load_nvfp4_gate() else: # hash if hash_lut is None: raise ValueError("hash router needs hash_lut") assert hash_lut.shape == (self.vocab_size, self.top_k), \ f"hash_lut shape {tuple(hash_lut.shape)} != " \ f"{(self.vocab_size, self.top_k)}" assert (hash_lut >= 0).all() and (hash_lut < self.num_experts).all(), \ "hash_lut contains out-of-range expert IDs" self.hash_lut = hash_lut.to(device=self.device, dtype=torch.int32) def load_nvfp4_gate(self, gate_lin) -> None: """Set the NVFP4 gate linear layer (2-kernel path). Called by the single_shot after constructing the Nvfp4Linear from checkpoint NVFP4 scales. When set, _run_dense_impl uses the production NVFP4 GEMM path instead of BF16 cuBLAS. """ self.gate_lin = gate_lin def load_nvfp4_fused_gate(self, gate_weight, gate_weight_scale, gate_ws2, gate_input_scale) -> None: """Set raw NVFP4 gate tensors for the fused single-kernel path. Preferred over load_nvfp4_gate (2-kernel) when available. The fused kernel handles activation quantization + GEMM + router epilogue in a single kernel launch. """ self.gate_weight = gate_weight.to(device=self.device) self.gate_weight_scale = gate_weight_scale.to(device=self.device) self.gate_ws2 = gate_ws2.to(device=self.device) if gate_ws2 is not None else None self.gate_input_scale = gate_input_scale.to(device=self.device) def finalize_weights(self) -> None: """Allocate output buffers and JIT-compile the routing kernel. Mirrors the finalize_weights() pattern in Nvfp4Linear: a one-time setup step called after all parameters are loaded. Triggers kernel compilation so the first forward isn't paying that cost. """ self._topk_weights_buf = torch.empty( self.max_num_tokens, self.top_k, dtype=torch.float32, device=self.device, ) self._topk_ids_buf = torch.empty( self.max_num_tokens, self.top_k, dtype=torch.int32, device=self.device, ) # Eager JIT — dispatcher knows our mode and triggers the right # kernel's compile path. See dsv4/ops/router.py. from dsv4.ops.router import warmup_router_compilation warmup_router_compilation(self) # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def __call__( self, hidden_states: torch.Tensor, token_ids: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Produce (topk_weights, topk_ids) for downstream Nvfp4MoE. Parameters ---------- hidden_states : Tensor [N, hidden_size] bfloat16 Required for dense mode. Ignored for hash mode (kept in the signature so the call site is mode-agnostic). token_ids : Tensor [N] int32, optional Required for hash mode. Ignored for dense mode. Returns ------- topk_weights : Tensor [N, top_k] float32 topk_ids : Tensor [N, top_k] int32 Notes ----- Both outputs are views into pre-allocated buffers — do not retain them across router calls. Nvfp4MoE consumes them immediately, which matches its existing contract. """ if self._topk_weights_buf is None: raise RuntimeError("Router.finalize_weights() not called") if self.mode == "dense": if hidden_states is None: raise ValueError("dense router requires hidden_states") return self._run_dense(hidden_states) else: if token_ids is None: raise ValueError("hash router requires token_ids") return self._run_hash(token_ids) # ------------------------------------------------------------------ # Mode-specific dispatch — each routes through a torch.library.custom_op # so Dynamo / torch.compile treats the kernel as opaque. # ------------------------------------------------------------------ def _run_dense(self, hidden_states: torch.Tensor): if self._runner_id is None: self._runner_id = register_router(self) return dense_router_op( hidden_states, self._runner_id, self.num_experts, self.top_k, ) def _run_hash(self, token_ids: torch.Tensor): if self._runner_id is None: self._runner_id = register_router(self) return hash_router_op( token_ids, self._runner_id, self.top_k, ) # ------------------------------------------------------------------ # Called by the custom_op dispatch in dsv4/ops/router.py — not by user code. # ------------------------------------------------------------------ def _run_dense_impl(self, hidden_states: torch.Tensor): """Hot-path: fused NVFP4, 2-kernel NVFP4, or BF16 fallback. Priority: 1. Fused NVFP4 kernel (single-kernel GEMM + router epilogue) 2. 2-kernel NVFP4 path (Nvfp4Linear + activation_topk) 3. BF16 cuBLAS fallback """ N = hidden_states.shape[0] out_w = self._topk_weights_buf[:N] out_ids = self._topk_ids_buf[:N] if self.gate_weight is not None: # Fused single-kernel path (preferred) from dsv4.kernels.router import dense_router_dispatch_nvfp4_fused dense_router_dispatch_nvfp4_fused( hidden_states=hidden_states, gate_weight=self.gate_weight, gate_weight_scale=self.gate_weight_scale, gate_ws2=self.gate_ws2, gate_input_scale=self.gate_input_scale, e_bias=self.e_bias, routed_scaling_factor=self.routed_scaling_factor, top_k=self.top_k, out_weights=out_w, out_ids=out_ids, ) elif self.gate_lin is not None: from dsv4.kernels.router import dense_router_dispatch_nvfp4 dense_router_dispatch_nvfp4( hidden_states=hidden_states, gate_lin=self.gate_lin, e_bias=self.e_bias, routed_scaling_factor=self.routed_scaling_factor, top_k=self.top_k, out_weights=out_w, out_ids=out_ids, ) else: from dsv4.kernels.router import dense_router_dispatch dense_router_dispatch( hidden_states=hidden_states, W_gate=self.W_gate, e_bias=self.e_bias, routed_scaling_factor=self.routed_scaling_factor, top_k=self.top_k, out_weights=out_w, out_ids=out_ids, ) return out_w, out_ids def _run_hash_impl(self, token_ids: torch.Tensor): """Hot-path entry into the hash gather kernel. Implementation lives in dsv4/kernels/cuda/hash_router.cu via the wrapper in dsv4/ops/router.py. """ from dsv4.kernels.router import hash_router_dispatch N = token_ids.shape[0] out_w = self._topk_weights_buf[:N] out_ids = self._topk_ids_buf[:N] hash_router_dispatch( token_ids=token_ids, hash_lut=self.hash_lut, top_k=self.top_k, out_weights=out_w, # filled with 1/k out_ids=out_ids, ) return out_w, out_ids