Gemma-4-26B-A4B-HybridAgent-4bit (TensorRT & Register-Fused Triton W4A16 — RTX 4090 24GB Compatible)

This repository contains the 4-bit quantized (compressed-tensors pack-quantized W4A16) edition of milkshake721/Gemma-4-26B-A4B-HybridAgent, designed to run on a single 24 GB GPU (e.g., RTX 4090 / 3090) (19.00 GB allocated VRAM) with TensorRT (torch_tensorrt) + Register-Fused Triton W4A16 MoE, preserving 100% of <latent> continuous hidden-state reasoning and <python> tool execution.

---

⚡ Benchmark Performance (4-bit W4A16 vs. BF16 with TensorRT)

Evaluated on the 10-question math & logic benchmark (eval_gemma4_26b_hybrid.py) on NVIDIA RTX PRO 6000 Blackwell (see BENCHMARK_AND_OPTIMIZATION_NOTES.md for full details):

Configuration	Allocated VRAM	Prefill (180 tok)	Single-Token Decode	10-Q Total Time	Avg Time / Q	10-Q Accuracy
BF16 Full (backend="tensorrt")	54.73 GB	313.03 ms	26.01 ms (38.5 tok/s)	106.80s	10.68s	10/10 (100%)
🏆 4-bit W4A16 (backend="tensorrt", Fused Triton W4A16 + F.rms_norm + Compiled FFN)	19.00 GB	187.51 ms	16.62 ms (60.2 tok/s)	83.65s	8.36s	10/10 (100%)

---

🔬 Quantization & Acceleration Architecture

1. 30 Layers $\times$ 128 MoE Routed Experts (experts.{0..127}.{gate,up,down}_proj):
   • Quantized to symmetric 4-bit (pack-quantized INT4 packed in int32, group_size=32, MSE-optimal scale search), shrinking MoE expert weights from 45.68 GB $\rightarrow$ 12.85 GB (19.00 GB total VRAM).
   • Register-Fused Triton W4A16 MoE Kernel (_fused_w4a16_moe_kernel): Unpacks INT4 weights, applies exact bfloat16 group-32 scales, and computes the dot product with activation $x$ entirely inside GPU registers/SRAM without ever materializing unpacked BF16 weights to HBM (2.02 ms / token across all 30 MoE layers, 2.28x faster than BF16).
2. Attention (self_attn), Router, Shared MLP, Embeddings, and HybridLink (hybridlink.pt):
   • Preserved in exact BF16 so <latent> entry/exit anchor transitions (token_id=262144 / 262145) and embed_scale = sqrt(2816) continuous recurrence suffer zero accuracy loss.
3. TensorRT HybridLink + Compiled FFN/RMSNorm Graph (enable_tensorrt_acceleration):
   • Compiles HybridLink (hybridlink.pt) into a native NVIDIA TensorRT ICudaEngine (torch_tensorrt.compile).
   • Replaces all 330 Gemma4RMSNorm calls per step with C++ F.rms_norm and compiles the post-attention FFN + Shared MLP + Router + Triton 4-bit MoE block across all 30 layers into a single fused graph (60.2 tok/s).

---

🚀 Quick Start 1: TensorRT Mode (backend="tensorrt", Recommended)

CLI Runner (sglang_tensorrt_runner.py):

python3 sglang_tensorrt_runner.py \
  --backend tensorrt \
  --model milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit \
  --question "A rectangle has a perimeter of 30 cm and an area of 54 cm^2. What are the lengths of its sides?"

Python API (load_hybrid_agent + HybridEngine):

import torch
from hybrid_agent.model import load_hybrid_agent
from hybrid_agent.engine import HybridEngine
from hybrid_agent.prompts import SYSTEM_PROMPT

model, tokenizer = load_hybrid_agent(
    "milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit",
    dtype=torch.bfloat16,
    device="cuda",
    backend="tensorrt",
)
engine = HybridEngine(model, tokenizer, max_new_tokens=1024, max_latent_steps=64, temperature=0.0)
traj = engine.run("If 3x + 7 = 22, what is the value of 5x - 2?", system_prompt=SYSTEM_PROMPT)
print("Final Answer:", traj.final_answer)  # -> 23
print("Latent Steps:", traj.latent_steps)
print(traj.output_rendered)

Run Verification & 10-Question Benchmark:

# 2-question verification test
python3 test_gemma4_26b_hybrid.py milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit --backend tensorrt

# 10-question benchmark evaluation (10/10 = 100.0% in ~83.6s)
python3 eval_gemma4_26b_hybrid.py milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit eval_results_4bit.json tensorrt

---

🚀 Quick Start 2: SGLang Mode (--backend sglang / sglang.launch_server)

# Run via SGLangHybridEngine (with TensorRT HybridLink & <python> tool loop)
python3 sglang_tensorrt_runner.py \
  --backend sglang \
  --model milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit \
  --question "If 3x + 7 = 22, what is the value of 5x - 2?"

# Or start a standalone SGLang Server
python3 -m sglang.launch_server \
  --model-path milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit \
  --quantization compressed-tensors \
  --dtype bfloat16 \
  --port 30000 \
  --mem-fraction-static 0.80 \
  --trust-remote-code