Initialize LLM and tokenizer.
(args: argparse.Namespace)
| 317 | |
| 318 | |
| 319 | def initialize_llm(args: argparse.Namespace) -> Tuple[LLM, AutoTokenizer]: |
| 320 | """Initialize LLM and tokenizer.""" |
| 321 | logger.info(f"Initializing LLM with model: {args.model_path}") |
| 322 | |
| 323 | try: |
| 324 | # Configure KV cache |
| 325 | kv_cache_config = KvCacheConfig( |
| 326 | # sparse attention doesn't support KV cache reuse |
| 327 | enable_block_reuse=False, |
| 328 | free_gpu_memory_fraction=args.kv_cache_fraction, |
| 329 | tokens_per_block=args.tokens_per_block, |
| 330 | ) |
| 331 | |
| 332 | # Configure CUDA graph |
| 333 | cuda_graph_config = CudaGraphConfig( |
| 334 | batch_sizes=args.cuda_graph_batch_sizes, |
| 335 | enable_padding=args.cuda_graph_padding_enabled, |
| 336 | ) if args.use_cuda_graph else None |
| 337 | |
| 338 | # Configure sparse attention |
| 339 | if args.rocket_sparse: |
| 340 | # Configure RocketKV sparse attention |
| 341 | sparse_attention_config = RocketSparseAttentionConfig( |
| 342 | window_size=args.window_size, |
| 343 | kernel_size=args.kernel_size, |
| 344 | prompt_budget=args.token_budget, |
| 345 | topk=args.topk, |
| 346 | kt_cache_dtype=args.kt_cache_dtype, |
| 347 | ) |
| 348 | logger.info(f"Using RocketKV sparse attention") |
| 349 | else: |
| 350 | sparse_attention_config = None |
| 351 | logger.info("Using standard attention") |
| 352 | |
| 353 | # Initialize LLM |
| 354 | llm = LLM( |
| 355 | model=args.model_path, |
| 356 | backend=args.backend, |
| 357 | kv_cache_config=kv_cache_config, |
| 358 | max_batch_size=args.max_batch_size, |
| 359 | attn_backend=args.attention_backend, |
| 360 | sparse_attention_config=sparse_attention_config, |
| 361 | tensor_parallel_size=args.tp_size, |
| 362 | moe_expert_parallel_size=args.moe_ep_size, |
| 363 | enable_attention_dp=args.enable_attention_dp, |
| 364 | max_seq_len=args.max_seq_len, |
| 365 | max_num_tokens=args.max_num_tokens, |
| 366 | cuda_graph_config=cuda_graph_config, |
| 367 | print_iter_log=args.print_iter_log, |
| 368 | moe_config=MoeConfig(backend=args.moe_backend), |
| 369 | ) |
| 370 | |
| 371 | # Initialize tokenizer |
| 372 | tokenizer = AutoTokenizer.from_pretrained(args.model_path) |
| 373 | |
| 374 | logger.info("LLM and tokenizer initialized successfully") |
| 375 | |
| 376 | return llm, tokenizer |
no test coverage detected