| """Fixed BF16 label-row projection for Flash Next; no vocabulary-sized GEMM.""" |
| import os |
|
|
|
|
| def register(): |
| from vllm import ModelRegistry |
| ModelRegistry.register_model( |
| 'Qwen4ExpForConditionalGeneration', |
| 'jev_flash_linear:Qwen4ExpForConditionalGeneration', |
| ) |
|
|
|
|
| |
| def __getattr__(name): |
| if name != 'Qwen4ExpForConditionalGeneration': |
| raise AttributeError(name) |
| import torch |
| from safetensors.torch import load_file |
| from torch import nn |
| from vllm.logger import init_logger |
| from vllm.models.qwen3_8_flash_next.nvidia.model import ( |
| Qwen3_8FlashNextForConditionalGeneration as BaseModel, |
| ) |
| logger = init_logger('vllm.jev_flash_linear') |
|
|
| class Qwen4ExpForConditionalGeneration(BaseModel): |
| def __init__(self, *, vllm_config, prefix=''): |
| super().__init__(vllm_config=vllm_config, prefix=prefix) |
| tensors = load_file(os.environ['JEV_LINEAR_HEAD'], device='cpu') |
| weight = tensors['weight'] |
| ids = tensors['token_ids'] |
| self.choice_head = nn.Linear(weight.shape[1], weight.shape[0], bias=False, |
| dtype=torch.bfloat16) |
| self.choice_head.weight.data.copy_(weight) |
| self.register_buffer('choice_token_ids', ids.to(self.choice_head.weight.device), |
| persistent=False) |
| self.choice_vocab_size = vllm_config.model_config.hf_text_config.vocab_size |
| del self.language_model.lm_head |
| del self.language_model.logits_processor |
| if any(key.endswith('lm_head') for key, _ in self.named_modules()): |
| raise RuntimeError('A full vocabulary output head remains') |
| if self.choice_head.weight.device.type != 'meta': |
| torch.testing.assert_close(self.choice_head.weight.cpu(), weight, rtol=0, atol=0) |
| self._choice_head_checked = False |
| logger.info('JEV compact linear head shape=%s dtype=%s device=%s; original lm_head removed', |
| tuple(weight.shape), self.choice_head.weight.dtype, |
| self.choice_head.weight.device) |
|
|
| def load_weights(self, weights): |
| loaded = super().load_weights( |
| (key, value) for key, value in weights if key != 'lm_head.weight' |
| ) |
| if loaded is not None: |
| loaded.add('choice_head.weight') |
| return loaded |
|
|
| def compute_logits(self, hidden_states): |
| scores = self.choice_head(hidden_states) |
| |
| |
| logits = scores.new_full((*scores.shape[:-1], self.choice_vocab_size), float('-inf')) |
| logits.index_copy_(-1, self.choice_token_ids, scores) |
| if not self._choice_head_checked: |
| logger.info('JEV linear forward hidden=%s scores=%s shape=%s; finite label logits=%s', |
| hidden_states.dtype, scores.dtype, tuple(scores.shape), |
| bool(torch.isfinite(scores).all())) |
| self._choice_head_checked = True |
| return logits |
|
|
| globals()[name] = Qwen4ExpForConditionalGeneration |
| return Qwen4ExpForConditionalGeneration |
|
|
