Transformers custommodeling · unknown
publisher-documented, not executed source
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
repo_name = "nvidia/Nemotron-Flash-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)
print('Initializing generation state...')
generation_state = model.init_cuda_graph_generation(
max_new_tokens=max_new_tokens,
batch_size=1,
device='cuda',
)
prompt = input("User:")
prompt = "User: " + prompt + "\nAssistant:"
inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
print(f"Generating with CUDA graph acceleration...")
outputs = model.generate_with_cuda_graph(
input_ids=inputs["input_ids"],
generation_state=generation_state,
max_new_tokens=256,
temperature=0,
eos_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
print(f"Response: {response}")outputs = model.generate_with_cache(
input_ids=inputs["input_ids"],
max_new_tokens=256,
temperature=0,
eos_token_id=tokenizer.eos_token_id,
)Prerequisites: Card pins torch<=2.9.1,transformers<=4.56.2,flash-attn<=2.7.3 plus causal-conv1d, mamba-ssm, flash-linear-attention.
- Requires custom CUDA graph/init generation methods; generic generate cannot bepresumedcompatible.
- Hybrid architecture fixed state is additional to attention KV; source hardware not universal.