"""Reference local runtime: upstream Naive attention and native Blackwell expert GEMMs. Loads ModelOpt's standard packed checkpoint without changing the source architecture. Expert dispatch is eager PyTorch; the matrix multiplications execute as NVFP4 W4A4. """ import os from pathlib import Path ROOT = Path(__file__).resolve().parents[1] os.environ.setdefault("HF_MODULES_CACHE", str(ROOT / "artifacts/hf-modules")) import torch import torch.nn as nn import torch.nn.functional as F from accelerate import init_empty_weights, dispatch_model from accelerate.utils import set_module_tensor_to_device from flashinfer import nvfp4_quantize, nvfp4_block_scale_interleave, mm_fp4, SfLayout from safetensors import safe_open from transformers import AutoConfig, AutoModelForCausalLM from safe_device_transfer import install as install_safe_transfers class NVFP4Linear(nn.Module): def __init__(self, weight, weight_scale, weight_scale_2, input_scale, backend="cutlass"): super().__init__() assert weight.dtype == torch.uint8 assert weight_scale.dtype == torch.float8_e4m3fn assert weight_scale.shape == (weight.shape[0], weight.shape[1] // 8) assert torch.isfinite(weight_scale.float()).all() assert weight_scale_2.numel() == input_scale.numel() == 1 assert weight_scale_2.item() > 0 and input_scale.item() > 0 self.register_buffer("weight", weight.contiguous()) self.register_buffer("scales", nvfp4_block_scale_interleave(weight_scale.contiguous().view(torch.uint8))) self.register_buffer("input_scale_inv", input_scale.float().reciprocal().reshape(1)) self.register_buffer("alpha", (weight_scale_2.float() * input_scale.float()).reshape(1)) self.backend = backend def forward(self, states): shape = states.shape states = states.reshape(-1, shape[-1]).contiguous() # FlashInfer's quantizer calls .cuda() internally, so its current device # must follow the layer's placement in a multi-GPU model. with torch.cuda.device(states.device): packed, scales = nvfp4_quantize(states, self.input_scale_inv, sfLayout=SfLayout.layout_128x4, do_shuffle=False) output = mm_fp4(packed, self.weight.T, scales, self.scales, self.alpha, out_dtype=states.dtype, backend=self.backend) return output.reshape(*shape[:-1], self.weight.shape[0]) class NVFP4Expert(nn.Module): def __init__(self, tensors, backend): super().__init__() for suffix in ("weight_scale_2", "input_scale"): if not torch.equal(tensors[f"gate_proj.{suffix}"], tensors[f"up_proj.{suffix}"]): raise ValueError(f"Exported gate/up {suffix} must match before fusion") self.gate_up = NVFP4Linear( torch.cat([tensors["gate_proj.weight"], tensors["up_proj.weight"]]), torch.cat([tensors["gate_proj.weight_scale"], tensors["up_proj.weight_scale"]]), tensors["gate_proj.weight_scale_2"], tensors["gate_proj.input_scale"], backend) self.down = NVFP4Linear(tensors["down_proj.weight"], tensors["down_proj.weight_scale"], tensors["down_proj.weight_scale_2"], tensors["down_proj.input_scale"], backend) def forward(self, states): gate, up = self.gate_up(states).chunk(2, dim=-1) return self.down(F.silu(gate) * up) class NVFP4Experts(nn.Module): def __init__(self, experts): super().__init__() self.experts = nn.ModuleList(experts) def forward(self, hidden_states, top_k_index, top_k_weights): output = torch.zeros_like(hidden_states) for expert_id in top_k_index.unique().tolist(): tokens, slots = torch.where(top_k_index == expert_id) current = self.experts[expert_id](hidden_states[tokens]) weighted = current * top_k_weights[tokens, slots, None].to(current.dtype) output.index_add_(0, tokens, weighted) return output def load_model(checkpoint, devices=None, backend="cutlass"): install_safe_transfers() os.environ.setdefault("HF_MODULES_CACHE", str(ROOT / "artifacts/hf-modules")) checkpoint = Path(checkpoint).resolve() devices = list(range(torch.cuda.device_count())) if devices is None else devices if not devices: raise ValueError("At least one Blackwell GPU is required") for device in devices: assert torch.cuda.get_device_capability(device)[0] >= 10 config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True) if config.model_type != "naive_n05_flash": raise ValueError(f"Unexpected model type: {config.model_type}") if hasattr(config, "quantization_config"): del config.quantization_config config.use_cache = True with init_empty_weights(): model = AutoModelForCausalLM.from_config(config, trust_remote_code=True, dtype=torch.bfloat16) device_map = {"model.embed_tokens": devices[0], "model.norm": devices[-1], "lm_head": devices[-1]} for i in range(config.num_hidden_layers): device_map[f"model.layers.{i}"] = devices[min(i * len(devices) // config.num_hidden_layers, len(devices) - 1)] import json index = json.loads((checkpoint / "model.safetensors.index.json").read_text())["weight_map"] loaded = set() # Each ModelOpt layer export is a self-contained shard. Expert tensors are # read by exact name so a missing scale fails instead of silently initializing. for layer_id, layer in enumerate(model.model.layers): prefix = f"model.layers.{layer_id}." gpu = device_map[prefix[:-1]] names = [name for name in index if name.startswith(prefix)] shards = {index[name] for name in names} if len(shards) != 1: raise ValueError("This loader requires one self-contained shard per decoder layer") with torch.cuda.device(gpu), safe_open(checkpoint / next(iter(shards)), framework="pt", device=gpu) as reader: if config.moe_layer_freq[layer_id]: if os.environ.get('NAIVE_EXPERT_BACKEND', 'reference') == 'fused': from fused_experts import FusedNVFP4Experts layer.mlp.experts = FusedNVFP4Experts(reader, f'{prefix}mlp.experts.', config.n_routed_experts, config.hidden_size, config.moe_intermediate_size, torch.device(f'cuda:{gpu}'), top_k=config.num_experts_per_tok, max_tokens=int(os.environ.get('PREFILL_CHUNK', '256'))) loaded.update(name for name in names if name.startswith(f'{prefix}mlp.experts.')) else: experts = [] for expert_id in range(config.n_routed_experts): ep = f"{prefix}mlp.experts.{expert_id}." tensors = {projection + "." + kind: reader.get_tensor(ep + projection + "." + kind) for projection in ("gate_proj", "up_proj", "down_proj") for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale")} experts.append(NVFP4Expert(tensors, backend)) loaded.update(ep + name for name in tensors) layer.mlp.experts = NVFP4Experts(experts) for name in names: if name in loaded: continue tensor = reader.get_tensor(name) set_module_tensor_to_device(model, name, gpu, value=tensor, dtype=tensor.dtype) loaded.add(name) print(f"Loaded layer {layer_id + 1}/{config.num_hidden_layers} on cuda:{gpu}", flush=True) remaining = set(index) - loaded for name in sorted(remaining): module_name = name.rsplit(".", 1)[0] if module_name not in device_map: raise ValueError(f"Unrecognized tensor outside decoder layers: {name}") gpu = device_map[module_name] with safe_open(checkpoint / index[name], framework="pt", device=gpu) as reader: tensor = reader.get_tensor(name) set_module_tensor_to_device(model, name, gpu, value=tensor, dtype=tensor.dtype) loaded.add(name) assert loaded == set(index) assert not any(p.is_meta for p in model.parameters()) model = dispatch_model(model, device_map=device_map, force_hooks=True) return model.eval().requires_grad_(False)