"""Laguna hybrid NVFP4 + calibrated EXL3 tail runtime for TP1 or TP2. Enabled only when LAGUNA_HYBRID_TIER is TR2 or TR3. The stock compressed- tensors NVFP4 MoE method is replaced with a fail-closed two-tier method: * saliency-selected hot experts are loaded byte-for-byte from the source checkpoint into a compact native vLLM CUTLASS NVFP4 kernel; * the remaining experts are loaded from the calibrated TP2 rank-sliced EXL3 tail; TP1 reconstructs the full expert by executing both stored slices into the kernel's additive fp32 output buffer; * both tiers consume the same router weights and their local outputs are added before vLLM performs its normal TP reduction. The trellis path uses preallocated scratch and fixed chunking and is therefore compatible with normal CUDA graph capture. This module does not change vLLM's execution mode or graph configuration. """ from __future__ import annotations import dataclasses import json import os import re import threading from pathlib import Path from typing import Any TIER = os.environ.get("LAGUNA_HYBRID_TIER", "").upper() if TIER in ("TR2", "TR3"): import torch import torch.nn as nn from safetensors import safe_open from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.fused_moe import SharedExperts from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( convert_to_nvfp4_moe_kernel_format, make_nvfp4_moe_kernel, make_nvfp4_moe_quant_config, select_nvfp4_moe_backend, ) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( CompressedTensorsW4A4Nvfp4MoEMethod as _StockNvfp4MoE, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Dynamic, kNvfp4Static, ) from vllm.model_executor.utils import set_weight_attrs from exllamav3.ext import exllamav3_ext as _exl3 BITS = 2 if TIER == "TR2" else 3 ARTIFACT_TP = 2 EXPERTS = 256 HIDDEN = 3072 INTERMEDIATE = 1024 TOP_K = 10 MCG = 0xCBAC1FED CHUNK = int(os.environ.get("LAGUNA_HYBRID_TRELLIS_CHUNK", "128")) if CHUNK <= 0: raise RuntimeError("LAGUNA_HYBRID_TRELLIS_CHUNK must be positive") MODEL_DIR = Path(os.environ.get("LAGUNA_HYBRID_MODEL_DIR", "/model")) TAIL_DIR = Path(os.environ.get("LAGUNA_HYBRID_TAIL_DIR", "/tail")) MAP_PATH = Path( os.environ.get( "LAGUNA_HYBRID_TIER_MAP", str(TAIL_DIR / f"laguna-{TIER.lower()}-tier-map.json"), ) ) if not MAP_PATH.is_file(): raise FileNotFoundError(f"Laguna hybrid tier map is missing: {MAP_PATH}") _MAP_PAYLOAD = json.loads(MAP_PATH.read_text(encoding="utf-8")) if _MAP_PAYLOAD.get("tier") != TIER: raise RuntimeError( f"tier map says {_MAP_PAYLOAD.get('tier')!r}, runtime requested {TIER}" ) if not _MAP_PAYLOAD.get("coverage_gate", {}).get("passed"): raise RuntimeError("Laguna hybrid tier map failed its coverage gate") _EXPECTED_BUDGET = { "TR2": { "tail_encoding": "trellis2", "tail_weight_bits": 2, "hot_experts_per_sparse_layer": 96, "tail_experts_per_sparse_layer": 160, "average_weight_bits": 2.75, }, "TR3": { "tail_encoding": "trellis3", "tail_weight_bits": 3, "hot_experts_per_sparse_layer": 64, "tail_experts_per_sparse_layer": 192, "average_weight_bits": 3.25, }, }[TIER] _BUDGET = _MAP_PAYLOAD.get("bit_budget") if not isinstance(_BUDGET, dict): raise RuntimeError("Laguna hybrid tier map lacks a bit-budget contract") _REQUIRED_BUDGET = { "hot_encoding": "nvfp4", "hot_weight_bits": 4, "tail_encoding": _EXPECTED_BUDGET["tail_encoding"], "tail_weight_bits": _EXPECTED_BUDGET["tail_weight_bits"], "hot_experts_per_sparse_layer": _EXPECTED_BUDGET[ "hot_experts_per_sparse_layer" ], "tail_experts_per_sparse_layer": _EXPECTED_BUDGET[ "tail_experts_per_sparse_layer" ], } for _field, _value in _REQUIRED_BUDGET.items(): if _BUDGET.get(_field) != _value: raise RuntimeError( f"Laguna {TIER} bit budget {_field}=" f"{_BUDGET.get(_field)!r} != {_value!r}" ) _average = _EXPECTED_BUDGET["average_weight_bits"] if ( float(_BUDGET.get("target_average_weight_bits", -1.0)) != _average or float(_BUDGET.get("nominal_average_weight_bits", -1.0)) != _average or _BUDGET.get("budget_passed") is not True ): raise RuntimeError( f"Laguna {TIER} tier map failed its exact bit-budget contract" ) _LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)") _RUNTIME: dict[tuple[int, int], dict[str, torch.Tensor | int]] = {} _RUNTIME_LOCK = threading.Lock() def _layer_index(value: str | None) -> int: match = _LAYER_RE.search(str(value or "")) if match is None: raise RuntimeError(f"cannot resolve Laguna layer index from {value!r}") layer = int(match.group(1)) if not 1 <= layer <= 47: raise RuntimeError(f"unexpected Laguna sparse layer index: {layer}") return layer def _tier_layer(layer: int) -> tuple[list[int], list[int]]: entry = _MAP_PAYLOAD["layers"].get(str(layer)) if entry is None: raise RuntimeError(f"tier map has no sparse layer {layer}") hot = [int(value) for value in entry["hot_experts"]] tail = [int(value) for value in entry["tail_experts"]] expected_hot = 96 if TIER == "TR2" else 64 expected_tail = EXPERTS - expected_hot if len(hot) != expected_hot or len(tail) != expected_tail: raise RuntimeError( f"layer {layer}: hot/tail counts {len(hot)}/{len(tail)} " f"!= {expected_hot}/{expected_tail}" ) if set(hot) | set(tail) != set(range(EXPERTS)) or set(hot) & set(tail): raise RuntimeError(f"layer {layer}: tier map is not a 256-expert partition") expected_tail_encoding = "trellis2" if TIER == "TR2" else "trellis3" if ( entry.get("hot_encoding") != "nvfp4" or entry.get("tail_encoding") != expected_tail_encoding ): raise RuntimeError( f"layer {layer}: tier encodings do not match {TIER}" ) return hot, tail def _runtime(device: torch.device, max_rows: int) -> dict[str, Any]: key = (device.index if device.index is not None else torch.cuda.current_device(), max_rows) with _RUNTIME_LOCK: existing = _RUNTIME.get(key) if existing is not None: return existing concurrency = int(_exl3.exl3_moe_max_concurrency(key[0])) value: dict[str, Any] = { "max_rows": max_rows, "cap": CHUNK, "xh": torch.empty((max_rows, HIDDEN), dtype=torch.float16, device=device), "out32": torch.empty((max_rows, HIDDEN), dtype=torch.float32, device=device), "tg": torch.empty( (concurrency, CHUNK, HIDDEN), dtype=torch.float16, device=device ), "tu": torch.empty( (concurrency, CHUNK, HIDDEN), dtype=torch.float16, device=device ), "ig": torch.empty( (concurrency, CHUNK, INTERMEDIATE // ARTIFACT_TP), dtype=torch.float16, device=device, ), "iu": torch.empty( (concurrency, CHUNK, INTERMEDIATE // ARTIFACT_TP), dtype=torch.float16, device=device, ), "flat_token": torch.arange( CHUNK, dtype=torch.int64, device=device ).repeat_interleave(TOP_K), "ones": torch.ones(CHUNK * TOP_K, dtype=torch.int64, device=device), } _RUNTIME[key] = value print( f"[laguna-hybrid] shared EXL3 runtime allocated: max_rows={max_rows} " f"chunk={CHUNK} concurrency={concurrency}", flush=True, ) return value class LagunaHybridNvfp4Exl3MoE(_StockNvfp4MoE): def __init__(self, moe, layer_name: str | None = None, use_a16: bool = False): super().__init__(moe, layer_name, use_a16) self.layer_name = layer_name self.layer_index = _layer_index(layer_name) self.hot, self.tail = _tier_layer(self.layer_index) self.hot_pos = {expert: index for index, expert in enumerate(self.hot)} self.tail_pos = {expert: index for index, expert in enumerate(self.tail)} self._seen: set[tuple[int, str, str]] = set() self.hot_kernel = None self.hot_layer = None self.hot_quant_config = None self.tail_slabs: dict[int, dict[str, dict[str, torch.Tensor]]] = {} self.tail_ptrs: dict[int, list[torch.Tensor]] = {} self.runtime_tp = 0 self.hot_expert_map = None self.tail_lut = None self.runtime = None def create_weights( self, layer: torch.nn.Module, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: if num_experts != EXPERTS or hidden_size != HIDDEN: raise RuntimeError( f"unexpected Laguna MoE geometry E={num_experts} H={hidden_size}" ) runtime_tp = get_tensor_model_parallel_world_size() if runtime_tp not in (1, ARTIFACT_TP): raise RuntimeError( "Laguna hybrid runtime supports tensor parallel world size 1 or 2" ) expected_local_i = INTERMEDIATE // runtime_tp if intermediate_size_per_partition != expected_local_i: raise RuntimeError( "Laguna hybrid intermediate partition mismatch: " f"{intermediate_size_per_partition} != {expected_local_i} " f"for TP{runtime_tp}" ) self.runtime_tp = runtime_tp layer.num_experts = num_experts layer.params_dtype = params_dtype hot_count = len(self.hot) rank = get_tensor_model_parallel_rank() def weight_loader( param, loaded, name_mapped=None, *, shard_id=None, expert_id=None, return_success=False, **_kwargs, ): expert = int(expert_id) if expert in self.tail_pos: return True if return_success else None local = self.hot_pos.get(expert) if local is None: return False if return_success else None name = str(name_mapped or "") shard = str(shard_id) if shard not in ("w1", "w2", "w3"): raise RuntimeError(f"unexpected expert shard {shard!r}") family = "w13" if ".w13_" in name else "w2" if "input_global_scale" in name: field = "input_global_scale" elif "weight_global_scale" in name: field = "weight_global_scale" elif "weight_scale" in name: field = "weight_scale" elif "weight_packed" in name: field = "weight_packed" else: raise RuntimeError(f"unrecognized compact NVFP4 parameter: {name}") if loaded.ndim >= 2 and runtime_tp == ARTIFACT_TP: if shard in ("w1", "w3"): loaded = loaded.chunk(ARTIFACT_TP, 0)[rank] else: loaded = loaded.chunk(ARTIFACT_TP, 1)[rank] destination = param.data[local] if family == "w13": if field in ("weight_packed", "weight_scale"): half = destination.shape[0] // 2 destination = ( destination[:half] if shard == "w1" else destination[half:] ) elif field in ("weight_global_scale", "input_global_scale"): destination = destination[0 if shard == "w1" else 1] destination.copy_(loaded.reshape(destination.shape).to(destination.dtype)) self._seen.add((expert, shard, field)) return True if return_success else None def parameter(name: str, shape: tuple[int, ...], dtype: torch.dtype) -> None: value = nn.Parameter( torch.empty( shape, dtype=dtype, device=torch.cuda.current_device(), ), requires_grad=False, ) set_weight_attrs(value, {**extra_weight_attrs, "weight_loader": weight_loader}) layer.register_parameter(name, value) local_i = expected_local_i parameter( "w13_weight_packed", (hot_count, 2 * local_i, HIDDEN // 2), torch.uint8, ) parameter( "w2_weight_packed", (hot_count, HIDDEN, local_i // 2), torch.uint8, ) parameter( "w13_weight_scale", (hot_count, 2 * local_i, HIDDEN // 16), torch.float8_e4m3fn, ) parameter( "w2_weight_scale", (hot_count, HIDDEN, local_i // 16), torch.float8_e4m3fn, ) parameter("w13_weight_global_scale", (hot_count, 2), torch.float32) parameter("w2_weight_global_scale", (hot_count,), torch.float32) parameter("w13_input_global_scale", (hot_count, 2), torch.float32) parameter("w2_input_global_scale", (hot_count,), torch.float32) print( f"[laguna-hybrid] layer {self.layer_index}: allocated " f"{hot_count} NVFP4 + {len(self.tail)} trellis{BITS} experts " f"for TP{runtime_tp} rank {rank}", flush=True, ) def _validate_hot_load(self) -> None: required = { (expert, shard, field) for expert in self.hot for shard in ("w1", "w2", "w3") for field in ( "weight_packed", "weight_scale", "weight_global_scale", "input_global_scale", ) } missing = required - self._seen if missing: raise RuntimeError( f"layer {self.layer_index}: missing {len(missing)} hot NVFP4 tensors; " f"first={sorted(missing)[:5]}" ) def _build_hot_kernel(self, layer) -> None: backend, experts_cls = select_nvfp4_moe_backend( config=dataclasses.replace( self.moe, num_experts=len(self.hot), num_local_experts=len(self.hot), num_logical_experts=len(self.hot), intermediate_size=self.moe.intermediate_size_per_partition, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), ), weight_key=kNvfp4Static, activation_key=kNvfp4Dynamic, ) kept_moe = dataclasses.replace( self.moe, num_experts=len(self.hot), num_local_experts=len(self.hot), num_logical_experts=len(self.hot), intermediate_size=self.moe.intermediate_size_per_partition, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), ) compact = nn.Module() compact.activation = getattr(layer, "activation", MoEActivation.SILU) compact.moe_config = kept_moe compact.local_num_experts = len(self.hot) compact.swiglu_limit = getattr(layer, "swiglu_limit", None) converted = convert_to_nvfp4_moe_kernel_format( nvfp4_backend=backend, layer=compact, w13=layer.w13_weight_packed, w13_scale=layer.w13_weight_scale, w13_scale_2=(1.0 / layer.w13_weight_global_scale[:, 0].contiguous()), a13_scale=(1.0 / layer.w13_input_global_scale), w2=layer.w2_weight_packed, w2_scale=layer.w2_weight_scale, w2_scale_2=(1.0 / layer.w2_weight_global_scale), a2_scale=(1.0 / layer.w2_input_global_scale), is_act_and_mul=True, ) ( compact.w13_weight, compact.w13_weight_scale, compact.w13_weight_scale_2, compact.w13_input_scale, compact.w2_weight, compact.w2_weight_scale, compact.w2_weight_scale_2, compact.w2_input_scale, ) = converted quant_config = make_nvfp4_moe_quant_config( backend=backend, w13_scale=compact.w13_weight_scale, w2_scale=compact.w2_weight_scale, w13_scale_2=compact.w13_weight_scale_2, w2_scale_2=compact.w2_weight_scale_2, a13_scale=compact.w13_input_scale, a2_scale=compact.w2_input_scale, swiglu_limit=compact.swiglu_limit, layer=compact, ) kernel = make_nvfp4_moe_kernel( moe_quant_config=quant_config, moe_config=kept_moe, experts_cls=experts_cls, backend=backend, routing_tables=None, layer=compact, ) kernel.fused_experts.process_weights_after_loading(compact) self.hot_kernel = kernel self.hot_layer = compact self.hot_quant_config = quant_config self.moe_kernel = kernel self._hot_keepalive = converted device = compact.w13_weight.device # CutlassExpertsFp4 explicitly does not support expert_map. Its # apply implementation accepts that argument for the modular # kernel interface but does not forward it to run_cutlass_moe_fp4. # Keep a lookup table here and remap router IDs before invoking the # compact kernel instead. hot_expert_map = torch.full( (EXPERTS,), -1, dtype=torch.int32, device=device ) for expert, local in self.hot_pos.items(): hot_expert_map[expert] = local self.hot_expert_map = hot_expert_map for name in ( "w13_weight_packed", "w2_weight_packed", "w13_weight_scale", "w2_weight_scale", "w13_weight_global_scale", "w2_weight_global_scale", "w13_input_global_scale", "w2_input_global_scale", ): if hasattr(layer, name): delattr(layer, name) print( f"[laguna-hybrid] layer {self.layer_index}: native hot kernel built " f"with backend={backend} experts={len(self.hot)}", flush=True, ) def _load_tail(self, layer) -> None: runtime_tp = get_tensor_model_parallel_world_size() rank = get_tensor_model_parallel_rank() if runtime_tp != self.runtime_tp: raise RuntimeError( f"Laguna hybrid TP changed during load: {self.runtime_tp} -> {runtime_tp}" ) artifact_ranks = ( [rank] if runtime_tp == ARTIFACT_TP else list(range(ARTIFACT_TP)) ) device = self.hot_layer.w13_weight.device count = len(self.tail) shapes = { "gate_proj": { "trellis": ( count, HIDDEN // 16, (INTERMEDIATE // ARTIFACT_TP) // 16, 16 * BITS, ), "suh": (count, HIDDEN), "svh": (count, INTERMEDIATE // ARTIFACT_TP), }, "up_proj": { "trellis": ( count, HIDDEN // 16, (INTERMEDIATE // ARTIFACT_TP) // 16, 16 * BITS, ), "suh": (count, HIDDEN), "svh": (count, INTERMEDIATE // ARTIFACT_TP), }, "down_proj": { "trellis": ( count, (INTERMEDIATE // ARTIFACT_TP) // 16, HIDDEN // 16, 16 * BITS, ), "suh": (count, INTERMEDIATE // ARTIFACT_TP), "svh": (count, HIDDEN), }, } slabs_by_rank: dict[int, dict[str, dict[str, torch.Tensor]]] = {} for artifact_rank in artifact_ranks: slabs: dict[str, dict[str, torch.Tensor]] = {} for projection, fields in shapes.items(): slabs[projection] = { field: torch.empty( shape, dtype=( torch.int16 if field == "trellis" else torch.float16 ), device=device, ) for field, shape in fields.items() } slabs_by_rank[artifact_rank] = slabs for local, expert in enumerate(self.tail): path = ( TAIL_DIR / TIER.lower() / f"layer-{self.layer_index:02d}" / f"expert-{expert:03d}.safetensors" ) if not path.is_file(): raise FileNotFoundError(f"missing Laguna tail expert artifact: {path}") with safe_open(str(path), framework="pt", device="cpu") as handle: metadata = handle.metadata() or {} expected_meta = { "format": "exl3-trellis", "bits": str(BITS), "tp": str(ARTIFACT_TP), "layer": str(self.layer_index), "expert": str(expert), "mcg_multiplier": hex(MCG), "hessian": "routed-real-activations", } for key, expected in expected_meta.items(): if metadata.get(key) != expected: raise RuntimeError( f"{path}: metadata {key}={metadata.get(key)!r} != {expected!r}" ) for artifact_rank in artifact_ranks: for projection in ("gate_proj", "up_proj", "down_proj"): marker = handle.get_tensor( f"{projection}.rank{artifact_rank}.mcg" ) if (int(marker.item()) & 0xFFFFFFFF) != MCG: raise RuntimeError(f"{path}: wrong MCG marker") for field in ("trellis", "suh", "svh"): source = handle.get_tensor( f"{projection}.rank{artifact_rank}.{field}" ) destination = slabs_by_rank[artifact_rank][projection][ field ][local] if tuple(source.shape) != tuple(destination.shape): raise RuntimeError( f"{path}:{projection}.rank{artifact_rank}.{field} " f"shape {tuple(source.shape)} != " f"{tuple(destination.shape)}" ) destination.copy_(source.to(destination.dtype)) self.tail_slabs = slabs_by_rank pointers_by_rank: dict[int, list[torch.Tensor]] = {} for artifact_rank, slabs in slabs_by_rank.items(): pointers: list[torch.Tensor] = [] for projection in ("gate_proj", "up_proj", "down_proj"): for field in ("trellis", "suh", "svh"): slab = slabs[projection][field] step = slab.stride(0) * slab.element_size() pointers.append( torch.tensor( [ slab.data_ptr() + index * step for index in range(count) ], dtype=torch.int64, device=device, ) ) pointers_by_rank[artifact_rank] = pointers self.tail_ptrs = pointers_by_rank tail_lut = torch.full((EXPERTS,), count, dtype=torch.int64, device=device) for expert, local in self.tail_pos.items(): tail_lut[expert] = local self.tail_lut = tail_lut from vllm.config import get_current_vllm_config max_rows = int( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) self.runtime = _runtime(device, max_rows) print( f"[laguna-hybrid] layer {self.layer_index}: loaded {count} calibrated " f"trellis{BITS} experts from artifact ranks {artifact_ranks} " f"for TP{runtime_tp} rank {rank}", flush=True, ) def _validate_execution_contract(self, layer) -> None: if not bool(getattr(self.moe, "is_act_and_mul", False)): raise RuntimeError( "Laguna hybrid requires SwiGLU act-and-multiply experts" ) activation = getattr(layer, "activation", MoEActivation.SILU) if activation != MoEActivation.SILU: raise RuntimeError( f"Laguna hybrid requires SiLU experts, got {activation!r}" ) swiglu_limit = getattr(layer, "swiglu_limit", None) if swiglu_limit not in (None, 0, 0.0): raise RuntimeError( "Laguna hybrid EXL3 tail does not support a nonzero " f"SwiGLU limit, got {swiglu_limit!r}" ) if bool(getattr(layer, "apply_router_weight_on_input", False)): raise RuntimeError( "Laguna hybrid requires router weights to be applied to " "expert outputs" ) def process_weights_after_loading(self, layer) -> None: self._validate_execution_contract(layer) self._validate_hot_load() self._build_hot_kernel(layer) self._load_tail(layer) def get_fused_moe_quant_config(self, layer): if self.hot_quant_config is None: raise RuntimeError("Laguna hot quant config requested before initialization") return self.hot_quant_config def _apply_tail(self, x, topk_weights, topk_ids): runtime = self.runtime if runtime is None or self.tail_lut is None: raise RuntimeError("Laguna trellis runtime is not initialized") rows = int(x.shape[0]) if rows > int(runtime["max_rows"]): raise RuntimeError( f"Laguna trellis rows {rows} exceed planned capacity " f"{runtime['max_rows']}" ) xh = runtime["xh"][:rows] xh.copy_(x) out = runtime["out32"][:rows] out.zero_() local_ids = self.tail_lut[topk_ids.long()] weights = topk_weights.to(torch.float16) count = len(self.tail) cap = int(runtime["cap"]) for start in range(0, rows, cap): chunk_rows = min(cap, rows - start) flat = local_ids[start : start + chunk_rows].reshape(-1) order = torch.argsort(flat) token = runtime["flat_token"][: chunk_rows * TOP_K].index_select( 0, order ).contiguous() sorted_weights = weights[start : start + chunk_rows].reshape(-1).index_select( 0, order ).contiguous() counts = torch.zeros(count + 1, dtype=torch.int64, device=x.device) counts.scatter_add_(0, flat, runtime["ones"][: chunk_rows * TOP_K]) # The EXL3 fused down projection scatter-adds with atomicAdd. # In TP1, executing both physically stored TP2 artifact slices # into the same zeroed fp32 buffer reconstructs the exact # tensor-parallel sum without allocating inside graph replay. for artifact_rank in sorted(self.tail_ptrs): _exl3.exl3_moe( xh[start : start + chunk_rows], out[start : start + chunk_rows], counts, token, sorted_weights, runtime["tg"], runtime["tu"], runtime["ig"], runtime["iu"], 0, BITS, BITS, BITS, *self.tail_ptrs[artifact_rank], True, False, True, False, True, False, 0.0, ) return out def apply( self, layer, x, topk_weights, topk_ids, shared_experts: SharedExperts | None, shared_experts_input, ): if ( self.hot_kernel is None or self.hot_expert_map is None or self.hot_layer is None ): raise RuntimeError("Laguna hot kernel is not initialized") # The router produces original IDs in [0, 256). Remap hot routes # into the compact [0, hot_count) space. Tail routes are sent to # the valid placeholder expert 0 with an exact zero router weight; # this avoids unsupported/out-of-range IDs without renormalizing # the original top-k weights. local_hot_ids = self.hot_expert_map[topk_ids.long()] hot_route = local_hot_ids >= 0 local_hot_ids = torch.where( hot_route, local_hot_ids, torch.zeros_like(local_hot_ids) ) local_hot_weights = topk_weights * hot_route.to(topk_weights.dtype) hot = self.hot_kernel.apply( x, self.hot_layer.w13_weight, self.hot_layer.w2_weight, local_hot_weights, local_hot_ids, activation=self.hot_layer.activation, global_num_experts=len(self.hot), expert_map=None, apply_router_weight_on_input=layer.apply_router_weight_on_input, shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) tail = self._apply_tail(x, topk_weights, topk_ids) return torch.add(hot.float(), tail).to(x.dtype) import vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 as _target _target.CompressedTensorsW4A4Nvfp4MoEMethod = LagunaHybridNvfp4Exl3MoE print( f"[laguna-hybrid] installed {TIER} runtime: bits={BITS} " f"artifact_tp={ARTIFACT_TP} runtime_tp=1-or-2 " f"map={MAP_PATH} tail={TAIL_DIR} CUDA graphs unchanged", flush=True, )