File size: 3,438 Bytes
559c709 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """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',
)
# Keep registration lightweight: vLLM imports the model class only when needed.
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)
# Keep the existing token-ID API. This scatter costs no full-vocab
# projection; the only output matrix multiplication is H x 86.
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
|