Una guía de codificación de extremo a extremo para NVIDIA KVPress para inferencia LLM de contexto largo, compresión de caché KV y generación de memoria eficiente

En este tutorial, adoptamos un enfoque práctico y detallado para explorar KVPress de NVIDIA y comprender cómo puede hacer que la inferencia de modelos de lenguaje de contexto largo sea más eficiente. Comenzamos configurando el entorno completo, instalando las bibliotecas necesarias, cargando un modelo Instruct compacto y preparando un flujo de trabajo simple que se ejecuta en Colab y al mismo tiempo demuestra el valor real de la compresión de caché KV. A medida que avanzamos en la implementación, creamos un corpus sintético de contexto largo, definimos preguntas de extracción específicas y ejecutamos múltiples experimentos de inferencia para comparar directamente la generación estándar con diferentes estrategias de KVPress. Al final del tutorial, habremos desarrollado una intuición más sólida sobre cómo funciona en la práctica la optimización del contexto prolongado, cómo los diferentes métodos de impresión afectan el rendimiento y cómo este tipo de flujo de trabajo se puede adaptar para la recuperación del mundo real, el análisis de documentos y aplicaciones LLM sensibles a la memoria.

importar sistema operativo, sistema, subproceso, ajuste de texto, tiempo, gc, json, matemáticas, aleatorio, advertencias, inspeccionar advertencias.filterwarnings(“ignorar”) def ejecutar(cmd): print(“\n[RUN]”, ” “.join(cmd)) subproceso.check_call(cmd) ejecutar([sys.executable, “-m”, “pip”, “install”, “-q”, “–upgrade”, “pip”]) correr([sys.executable, “-m”, “pip”, “install”, “-q”, “torch”, “transformers”, “accelerate”, “bitsandbytes”, “sentencepiece”, “kvpress==0.4.0”]) intente: desde google.colab importe datos de usuario hf_token = userdata.get(“HF_TOKEN”) excepto excepción: hf_token = os.environ.get(“HF_TOKEN”, “”) si no es hf_token: intente: importe getpass hf_token = getpass.getpass(“Ingrese su token de cara de abrazo (déjelo vacío si el modelo es público y accesible): “).strip() excepto excepción: hf_token = “” si hf_token: os.environ[“HF_TOKEN”] = hf_token os.environ[“HUGGINGFACEHUB_API_TOKEN”] = hf_token importar antorcha importar transformadores importar kvpress desde la tubería de importación de transformadores, BitsAndBytesConfig desde kvpress importar ExpectedAttentionPress, KnormPress print(“Python:”, sys.version.split()[0]) print(“Antorcha:”, torch.__version__) print(“Transformers:”, transformadores.__version__) print(“CUDA disponible:”, torch.cuda.is_available()) if torch.cuda.is_available(): print(“GPU:”, torch.cuda.get_device_name(0)) MODEL_ID = “Qwen/Qwen2.5-1.5B-Instruct” MAX_NEW_TOKENS = 96 SEMILLA = 42 aleatoria.semilla(SEMILLA) antorcha.manual_semilla(SEMILLA)

Configuramos el entorno Colab e instalamos todas las bibliotecas necesarias para ejecutar el flujo de trabajo KVPress con éxito. Recopilamos de forma segura el token de Hugging Face, configuramos variables de entorno e importamos los módulos principales necesarios para la carga del modelo, la ejecución de la canalización y los experimentos de compresión. También imprimimos los detalles del tiempo de ejecución y del hardware para que entendamos claramente la configuración en la que realizamos el tutorial.

si torch.cuda.is_available(): torch.cuda.empty_cache() quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type=”nf4″, bnb_4bit_use_double_quant=True, ) pipe = pipeline( “kv-press-text-generación”, modelo=MODEL_ID, dispositivo_map=”auto”, token=hf_token si hf_token else Ninguno, model_kwargs={ “quantization_config”: quantization_config, “attn_implementation”: “sdpa”, }, ) más: pipe = pipeline( “kv-press-text-generación”, modelo=MODEL_ID, dispositivo_map=”auto”, torch_dtype=torch.float32, token=hf_token if hf_token else Ninguno, model_kwargs={ “attn_implementation”: “sdpa”, }, ) def cuda_mem(): si no, torch.cuda.is_available(): return {“allocated_gb”: Ninguno, “reserved_gb”: Ninguno, “peak_gb”: Ninguno} return { “allocated_gb”: round(torch.cuda.memory_allocated() / 1024**3, 3), “reserved_gb”: round(torch.cuda.memory_reserved() / 1024**3, 3), “peak_gb”: round(torch.cuda.max_memory_allocated() / 1024**3, 3), } def reset_peak(): si torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def extract_answer(x): si isinstance(x, list) y len(x) > 0: x = x[0]
si es instancia (x, dict): para k en [“answer”, “generated_text”, “text”, “output_text”]: si k en x: devolver x[k]
return json.dumps(x, sangría=2, asegurar_ascii=False) return str(x) def generate_once(contexto, pregunta, presione=Ninguno, etiqueta=”ejecutar”): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() reset_peak() start = time.time() out = pipe( contexto, pregunta=pregunta, presione=presione, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, temperatura=Ninguna, return_full_text=False,) transcurrido = time.time() – iniciar respuesta = extraer_respuesta(out) stats = cuda_mem() resultado = { “label”: etiqueta, “elapsed_sec”: round(elapsed, 2), “allocated_gb”: stats[“allocated_gb”]”reservado_gb”: estadísticas[“reserved_gb”]”peak_gb”: estadísticas[“peak_gb”]”respuesta”: respuesta.strip(), } devuelve resultado

Inicializamos la canalización kv-press-text-generación y la configuramos de manera diferente dependiendo de si hay compatibilidad con GPU disponible. Definimos las funciones auxiliares que miden el uso de la memoria CUDA, restablecen la memoria máxima, extraen respuestas de los resultados del modelo y ejecutan un pase de generación única de manera limpia. Esta parte proporciona la lógica de ejecución reutilizable que impulsa el resto del tutorial y nos permite comparar la inferencia de referencia con la compresión de caché KV.

registros_empresa = [
{“company”: “Arcturus Dynamics”, “hq”: “Bengaluru”, “founded”: 2017, “focus”: “warehouse robotics”},
{“company”: “BlueMesa Energy”, “hq”: “Muscat”, “founded”: 2014, “focus”: “grid analytics”},
{“company”: “CinderPeak Health”, “hq”: “Pune”, “founded”: 2019, “focus”: “clinical imaging AI”},
{“company”: “DeltaForge Marine”, “hq”: “Kochi”, “founded”: 2012, “focus”: “autonomous vessel telemetry”},
{“company”: “EonCircuit Labs”, “hq”: “Hyderabad”, “founded”: 2020, “focus”: “edge silicon tooling”},
{“company”: “Frostline Aero”, “hq”: “Jaipur”, “founded”: 2016, “focus”: “drone inspection”},
]

hechos_aguja = [
“PROJECT NEEDLE 1: The internal codename for the confidential pilot program is SAFFRON-17.”,
“PROJECT NEEDLE 2: The audit escalation owner is Meera Vashisht.”,
“PROJECT NEEDLE 3: The approved deployment region for the first production rollout is Oman North.”,
“PROJECT NEEDLE 4: The emergency rollback phrase is amber lantern.”,
“PROJECT NEEDLE 5: The signed commercial start date is 17 September 2026.”,
]

background_block = “”” Los sistemas de contexto largo a menudo contienen notas operativas repetidas, registros históricos, secciones de políticas y artefactos de recuperación ruidosos. El objetivo de esta demostración es crear un mensaje realistamente largo donde solo unos pocos detalles importan para la respuesta posterior. La compresión de caché KV reduce el uso de memoria al podar los pares clave-valor almacenados en caché y al mismo tiempo preserva la calidad de la respuesta. “”” Policy_block = “”” Resumen de la política operativa: 1. La seguridad anula el rendimiento cuando la confianza del sensor cae por debajo del umbral. 2. Los registros deben preservar la región, la marca de tiempo, la clase de dispositivo y el estado de aprobación del operador. 3. Las pruebas de campo pueden contener anexos duplicados, artefactos de estilo OCR y resúmenes de cumplimiento repetidos. 4. Un buen modelo de contexto largo debe ignorar las repeticiones irrelevantes y recuperar los detalles específicos que importan. []
para i en el rango (120): rec = company_records[i % len(company_records)]
records_text.append( f”Registro {i+1}: {rec[‘company’]} tiene su sede en {rec[‘hq’]}, fundada en {rec[‘founded’]}, y se centra en {rec[‘focus’]}. ” f”Memorando trimestral {i+1}: la retención se mantuvo estable, la capacitación de los operadores avanzó y el apéndice de cumplimiento se volvió a adjuntar para su revisión.” ) Needle_insert_positions = {18, 41, 73, 96, 111} full_corpus = []
para i, para en enumerate(records_text): full_corpus.append(background_block.strip()) full_corpus.append(policy_block.strip()) full_corpus.append(para) si estoy en Needle_insert_positions: full_corpus.append(needle_facts[len([x for x in needle_insert_positions if x <= i]) – 1])

Creamos un conjunto de datos sintéticos de contexto largo para probar el sistema KVPress de una manera controlada pero realista. Definimos registros de la empresa, insertamos hechos ocultos importantes en diferentes posiciones y los mezclamos con repetidos bloques de antecedentes y políticas, lo que hace que el mensaje sea largo y ruidoso. Esto nos ayuda a simular el contexto en el que la inferencia eficiente en la memoria es importante y el modelo debe recuperar solo los detalles verdaderamente relevantes.

contexto = “\n\n”.join(full_corpus) question = textwrap.dedent(“”” Responda usando solo el contexto proporcionado. Proporcione un objeto JSON compacto con exactamente estas claves: fecha_inicio_comercial región_despliegue propietario_auditoría frase_retroceso nombre_código_piloto “””).strip() print(“\nCaracteres de contexto:”, len(contexto)) print(“Palabras aproximadas:”, len(context.split())) experimentos = []

línea de base = generar_once (contexto, pregunta, presione = Ninguno, etiqueta = “línea de base_no_compresión”) experimentos.append (línea de base) presiones = [
(“expected_attention_0.7”, ExpectedAttentionPress(compression_ratio=0.7)),
(“expected_attention_0.5”, ExpectedAttentionPress(compression_ratio=0.5)),
(“knorm_0.5”, KnormPress(compression_ratio=0.5)),
]

para etiqueta, presione presiona: intente: resultado = generar_once(context, pregunta, presione=presione, etiqueta=etiqueta) experimentos.append(resultado) excepto Excepción como e: experimentos.append({ “label”: etiqueta, “elapsed_sec”: Ninguno, “allocated_gb”: Ninguno, “reserved_gb”: Ninguno, “peak_gb”: Ninguno, “answer”: f”FAILED: {type(e).__name__}: {e}” }) intente: desde kvpress import DecodingPress sig = inspect.signature(DecodingPress) kwargs = {“base_press”: KnormPress()} si “compression_interval” en sig.parameters: kwargs[“compression_interval”] = 10 elif “pasos_compresión” en sig.parameters: kwargs[“compression_steps”] = 10 si “target_size” en sig.parameters: kwargs[“target_size”] = 512 elif “token_buffer_size” en sig.parameters: kwargs[“token_buffer_size”] = 512 si “hidden_states_buffer_size” en sig.parameters: kwargs[“hidden_states_buffer_size”] = 0 decoding_press = DecodingPress(**kwargs) decoding_result = generate_once(context, question, press=decoding_press, label=”decoding_knorm”) experiments.append(decoding_result) excepto Excepción como e: experiments.append({ “label”: “decoding_knorm”, “elapsed_sec”: Ninguno, “allocated_gb”: Ninguno, “reserved_gb”: Ninguno, “peak_gb”: Ninguno, “respuesta”: f”SKIPPED_OR_FAILED: {tipo(e).__nombre__}: {e}” })

Reunimos el contexto final, definimos la pregunta de extracción estructurada y lanzamos el conjunto central de experimentos de inferencia. Primero ejecutamos la línea base sin compresión, luego aplicamos múltiples estrategias de prensa para observar cómo las diferentes relaciones de compresión afectan los resultados. También llevamos a cabo un experimento de compresión orientado a la decodificación, que extiende el tutorial más allá del precarga y proporciona una visión más amplia del marco KVPress.

print(“\n” + “=” * 120) print(“RESULTADOS”) print(“=” * 120) para r en experimentos: print(f”\n[{r[‘label’]}]”) print(“segundo_transcurrido:”, r[“elapsed_sec”]) print(“asignado_gb:”, r[“allocated_gb”]) print(“reservado_gb:”, r[“reserved_gb”]) imprimir(“pico_gb:”, r[“peak_gb”]) imprimir(“respuesta:”) imprimir(r[“answer”]) print(“\n” + “=” * 120) print(“RESUMEN SIMPLE”) print(“=” * 120) def safe_float(x): intente: devolver float(x) excepto Excepción: devolver Ninguna base_peak = safe_float(baseline[“peak_gb”]) si baseline.get(“peak_gb”) no es Ninguno más Ninguno base_time = safe_float(baseline[“elapsed_sec”]) si baseline.get(“elapsed_sec”) no es Ninguno más Ninguno para r en experimentos[1:]: pico = flotador_seguro(r[“peak_gb”]) t = flotador_seguro(r[“elapsed_sec”]) Peak_delta = Ninguno si base_peak es Ninguno o el pico es Ninguno más round(base_peak – pico, 3) time_delta = Ninguno si base_time es Ninguno o t es Ninguno más round(base_time – t, 2) print({ “label”: r[“label”]”peak_gb_saved_vs_baseline”: pico_delta, “time_sec_saved_vs_baseline”: time_delta, “answer_preview”: r[“answer”][:180].replace(“\n”, ” “) }) print(“\n” + “=” * 120) print(“SIGUIENTES PASOS OPCIONALES”) print(“=” * 120) print(“1. Cambie MODEL_ID por un modelo de instrucciones de contexto largo más potente que se ajuste a su GPU.”) print(“2. Aumente la longitud del contexto duplicando records_text más veces.”) print(“3. Pruebe otras prensas de kvpress, como SnapKVPress, StreamingLLMPress, QFilterPress o ChunkKVPress.”) print(“4. Reemplace el corpus sintético con sus propios fragmentos largos de PDF/texto y mantenga el mismo ciclo de evaluación.”)

Imprimimos todos los resultados del experimento en un formato legible y resumimos las diferencias de tiempo de ejecución y memoria en relación con la línea de base. Calculamos métricas de comparación simples para ver rápidamente cuánta memoria o tiempo ahorra cada estrategia de compresión. Luego concluimos con los siguientes pasos sugeridos para extender el tutorial a modelos más sólidos, contextos más largos, métodos de impresión adicionales y cargas de trabajo de documentos del mundo real.

En conclusión, desarrollamos una sólida comprensión práctica de cómo se puede utilizar KVPress de NVIDIA para optimizar la inferencia de contexto largo en un entorno realista basado en Colab. Hicimos más que simplemente ejecutar un modelo: creamos un flujo de trabajo de extremo a extremo que instala el marco, carga la canalización correctamente, construye una entrada significativa de contexto largo, aplica múltiples presiones de compresión y evalúa los resultados en términos de calidad de respuesta, tiempo de ejecución y comportamiento de la memoria. Al comparar la generación de referencia con la generación de caché KV comprimida, vimos claramente las compensaciones involucradas. Obtuvimos una intuición útil sobre cuándo estos métodos pueden ayudar a reducir la presión sobre los recursos sin dañar gravemente la fidelidad de la producción. También exploramos la flexibilidad del marco probando diferentes configuraciones de prensa e incluyendo una ruta de compresión opcional orientada a la decodificación, brindando una visión más amplia de cómo se puede usar KVPress más allá de un único ejemplo estático.

Consulta los Códigos y el Cuaderno aquí. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros