--- library_name: transformers pipeline_tag: text-generation language: - zh - en license: other license_name: baihu-custom-license license_link: https://huggingface.co/NovaAI6868/BaiHu-V1-Flash/blob/main/LICENSE.custom.md base_model: Qwen/Qwen3-0.6B-Base tags: - sparse-attention - subq - ssa - long-context - commercial-license-required - text-generation --- # BaiHu-V1-Flash **BaiHu-V1-Flash** is a retrofit of the dense-attention model `Qwen/Qwen3-0.6B-Base` into an **SSA (Sparse-attention + SubQ)** architecture, obtained by continued pretraining. - Base model: `Qwen/Qwen3-0.6B-Base` (28 layers / 16 Q heads / 8 KV heads / head_dim 128 / 32K context / tied embeddings) - Parameters: 598.8M - Training data: mixed Chinese + English (Fineweb-Edu-Chinese-V2.1 + fineweb-edu, 50/50) - License: **free for personal use; a paid license is required for commercial use** (see "License" below) --- ## 1. Architecture: SSA (three paths, each with its own softmax, then summed) Every layer keeps the base model's MLP / RMSNorm weights and replaces full attention with an SSA layer built from three parallel paths: | Path | Role | Complexity | |---|---|---| | `shared` | every query sees all **completed** blocks through one compressed vector per block | `O(T·T/B)` | | `local` | dense causal attention over the most recent window | `O(T·w)` | | `sparse` (**SubQ**) | only 4 of 16 query heads produce block scores, shared across the whole head group; real attention is computed only for the selected top-k blocks | `O(T·k·B)` | ### Hyperparameters | Parameter | Value | Meaning | |---|---|---| | `ssa_block_size` | 64 | block size B | | `ssa_top_k` | 8 | number of blocks selected by the sparse path | | `ssa_local_blocks` | 2 | local window = 3 × 64 = 192 tokens | | `ssa_num_subq_heads` | 4 | SubQ heads, r = 16 / 4 = 4 | | `ssa_router_dim` | 128 | router subspace dimension | | `ssa_compress_dim` | 128 | block compression dimension | Only **2.75M** parameters are new (≈0.46% of the model); all other weights are inherited from the base model. --- ## 2. Training | Item | Setting | |---|---| | Starting point | a conversion checkpoint that is **bit-exact** with the base model (`max\|Δlogit\| = 0.000e+00`) | | Tokens seen | 5.0M (≈0.25 epoch of the corpus) | | Sequence length | 512 (must be a multiple of `ssa_block_size = 64`) | | Effective batch | 8192 tokens (batch 2 × grad_accum 8) | | Precision | fp32 | | Optimizer | SGD with momentum 0.9 | | Learning rate | 5e-4, 50-step warmup, cosine decay to 10% | | Hardware | single NVIDIA GTX TITAN X (Maxwell, sm_52, 12.9 GB) | | Throughput | ≈274 tokens/s | ### Critical issues found and fixed during this retrofit Several defects silently break training and are worth documenting: 1. **Both new branches had identically zero gradients (blocking).** To make the converted model bit-exact with the base model, `compress_out` and `router_out` were initialized to exactly zero, and the branches were skipped entirely by a gate. The branch output was therefore always zero, so the back-propagated gradient was also always zero: all 2.75M SSA parameters **stayed frozen for the entire run** and the sparse attention was dead code. The fix is a small non-zero initialization. 2. **Routing was non-differentiable.** `top-k` produces hard indices, and indexing is not differentiable. If the routing scores are used only to decide *which* blocks to read and never enter the softmax, the gradients of `router_q` / `router_k` are **exactly zero** — the router can never learn to route. The fix is to feed the selected blocks' scores, squashed through `tanh` and gently scaled, into the attention logits as an additive bias. 3. **The shared summary was a sum, not a mean.** Its magnitude grew linearly with the prefix, and because that branch is injected at full weight (a single-element softmax has probability exactly 1), it swamped the residual stream: hidden states grew from 0.2 to about 7 in layer 0 and to about 1900 by layer 27, and validation loss went 3.56 → 11.38. The fix is to divide by the token count. 4. **The shared branch needs an explicit gate.** With a single-element softmax the probability is always 1, so the initialization scale of `compress_out` cannot control the injection strength at all (measured: scales from 1e-4 to 0.03 all left the loss at exactly 7.3526). A learnable scalar gate, initialized to a small positive value, lets the optimizer decide how far to open it. --- ## 3. How to Run Inference ### 3.1 Important: this is a custom architecture `BaiHu-V1-Flash` uses `model_type: baihu_ssa`, which is **not** in the Transformers registry. Loading it with a plain `AutoModelForCausalLM.from_pretrained(...)` fails with: ``` ValueError: The checkpoint you are trying to load has model type `baihu_ssa` but Transformers does not recognize this architecture. ``` You must register the config and model classes first. This is a one-time, three-line step (see below). Also note: **this model does not go through `transformers.GenerationMixin`.** SSA owns its own KV-cache layout (`BaiHuSSACache`) because the cache stores per-block compressed prefix sums rather than a plain growing key/value tensor. Call the model's own `generate()` method; beam search is not supported. ### 3.2 Setup ```bash # The SSA implementation lives in the project repository (not on the Hub), # because it is a custom architecture. git clone ssa_model cd ssa_model uv venv --python 3.11 .venv # Ampere or newer (RTX 30xx/40xx, A100, ...): any recent torch works. uv pip install --python .venv/Scripts/python.exe \ "numpy>=1.26" "transformers>=4.51" "safetensors>=0.4" "torch>=2.6" # Maxwell / Pascal / Volta (GTX 9xx/10xx, TITAN X, V100, ...): pin torch 2.7.1 — # see the GPU note below for why `torch>=2.6` would resolve to a build without your kernels. uv pip install --python .venv/Scripts/python.exe \ "numpy>=1.26" "transformers>=4.51" "safetensors>=0.4" "torch==2.7.1" ``` > **GPU note (Maxwell and older).** If you are on an NVIDIA Maxwell card such as the > GTX TITAN X (compute capability sm_52), recent PyTorch wheels no longer ship kernels for > it: PyTorch **2.8 removed sm_50/sm_60**, and 2.8.0+cu126 only contains `sm_61…sm_90`. > You will get `CUDA error: no kernel image is available for execution on the device`. > Use **torch 2.7.1+cu126**, which still contains `sm_50`. Also train/infer in **fp32** on > Maxwell: measured fp32 5.30 / fp16 4.46 / bf16 3.20 TFLOPS, so half precision is a loss. ### 3.3 Minimal working example ```python import os import sys import torch from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer # Make the custom implementation importable (path to the cloned repo's src/) sys.path.insert(0, os.path.join("ssa_model", "src")) from baihu_ssa.configuration_baihu_ssa import BaiHuSSAConfig from baihu_ssa.model_baihu_ssa import BaiHuSSAForCausalLM # ---- register the custom architecture (REQUIRED) ---- AutoConfig.register("baihu_ssa", BaiHuSSAConfig, exist_ok=True) AutoModelForCausalLM.register(BaiHuSSAConfig, BaiHuSSAForCausalLM, exist_ok=True) REPO = "NovaAI6868/BaiHu-V1-Flash" device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float32 if device == "cpu" else torch.float32 # see GPU note: fp32 tokenizer = AutoTokenizer.from_pretrained(REPO) model = AutoModelForCausalLM.from_pretrained(REPO, dtype=dtype).to(device).eval() prompt = "人工智能的未来是" input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) with torch.no_grad(): out = model.generate( input_ids, max_new_tokens=32, do_sample=False, # greedy; supported # do_sample=True, temperature=0.8, top_k=50, # sampling also supported ) print(tokenizer.decode(out[0], skip_special_tokens=True)) ``` ### 3.4 `generate()` parameters The built-in `generate()` is a self-contained decoding loop (greedy or temperature/top-k sampling). Supported arguments: | Argument | Default | Meaning | |---|---|---| | `max_new_tokens` | 32 | number of tokens to generate | | `do_sample` | `False` | `False` = greedy, `True` = sample | | `temperature` | 1.0 | sampling temperature (used when `do_sample=True`) | | `top_k` | `None` | top-k sampling cutoff (used when `do_sample=True`) | | `eos_token_id` | `None` | stop early when all sequences emit this token | | `use_cache` | `True` | keep the SSA cache between steps; **leave on**, decoding is much slower without it | **Not supported:** beam search (the SSA cache has no `reorder_cache`), and `transformers` generation utilities such as `logits_processor` / `stopping_criteria`. ### 3.5 Computing perplexity / loss ```python with torch.no_grad(): inputs = tokenizer("Your text here", return_tensors="pt").to(device) loss = model(**inputs, labels=inputs["input_ids"]).loss print("loss", loss.item(), "ppl", torch.exp(loss).item()) ``` ### 3.6 Throughput you should expect Measured on a GTX TITAN X (sm_52, 12.9 GB), fp32, prompt 1024 tokens, 64 generated tokens: | | Qwen3-0.6B-Base | BaiHu-V1-Flash | |---|---|---| | Prefill / TTFT | 503 ms | 1431 ms | | Per-token decode (TPOT) | 36.6 ms | 80.5 ms | | Decode throughput | 27.3 tok/s | 12.4 tok/s | | Peak memory (generation) | 9.91 GB | 9.01 GB | **BaiHu-V1-Flash is currently ~2.2× slower to decode than the base model**, even though the sparse path reads fewer keys. The reason is implementation, not architecture: the attention layer loops over query blocks in Python and launches many small kernels, so kernel-launch overhead dominates the FLOPs saved. Memory use is lower; speed is not yet a win. See "Known Limitations" below. ### 3.7 Expected output quality (honest warning) This checkpoint has seen only **5.0M tokens**, so the new branches are far from converged. Greedy decoding tends to **repeat itself** and it is **worse than the base model on perplexity** (ppl 82.3 vs 39.4 at 512 tokens). A real greedy sample: ``` prompt: 人工智能的未来是 output: 人工智能的未来是怎样的？人工智能的未来是怎样的？人工智能的未来是怎样的？… ``` Use it to study or continue the SSA retrofit — **do not** expect it to match `Qwen3-0.6B-Base` as a general-purpose model yet. ### 3.8 Troubleshooting | Error | Cause | Fix | |---|---|---| | `does not recognize this architecture` / `KeyError: 'baihu_ssa'` | custom architecture not registered | call `AutoConfig.register` + `AutoModelForCausalLM.register` as in 3.3 | | `CUDA error: no kernel image is available for execution on the device` | installed PyTorch has no kernels for your GPU (Maxwell/sm_52) | install torch 2.7.1+cu126 or older (2.8 dropped sm_50/sm_60) | | `AttributeError: ... has no attribute 'tie_weights'` | an external tool treating the model as `PreTrainedModel` | use a version of the implementation that defines `tie_weights()` (already fixed upstream) | | `BaiHuSSACache` errors when calling `model.generate(...)` from GenerationMixin | this model bypasses `GenerationMixin` | call `model.generate(...)` on the BaiHu model itself | --- ## 4. Evaluation ### 4.1 Language modeling perplexity (validation set, identical windows) | Sequence length | Qwen3-0.6B-Base | BaiHu-V1-Flash | |---|---|---| | 512 | 3.6726 / ppl 39.354 | 4.4108 / ppl 82.335 | | 1024 | 3.4691 / ppl 32.109 | 4.2575 / ppl 70.631 | | 2048 | 3.0801 / ppl 21.761 | 3.9213 / ppl 50.467 | ### 4.2 Standard benchmarks (lm-evaluation-harness) | Task | Qwen3-0.6B-Base | BaiHu-V1-Flash | Delta | |---|---|---|---| | arc_easy | 0.5550 | 0.6250 | +0.0700 | | hellaswag | 0.5350 | 0.5200 | -0.0150 | | piqa | 0.7050 | 0.6950 | -0.0100 | | winogrande | 0.6300 | 0.6300 | +0.0000 | ### 4.3 Inference compute and resource usage | Metric | Qwen3-0.6B-Base | BaiHu-V1-Flash | |---|---|---| | Parameters (M) | 596.0500 | 598.8000 | | Prefill peak memory (GB) | 9.6200 | 8.0210 | | Generation peak memory (GB) | 9.9120 | 9.0050 | | Prefill latency (s) | 0.5030 | 1.4310 | | TTFT (ms) | 502.7 | 1431.2 | | TPOT (ms) | 36.6 | 80.5 | | Decode throughput (tok/s) | 27.3300 | 12.4200 | | Attention FLOPs/token (GFLOPs) | 0.1176 | 0.1057 | | Attention key accesses vs full attention | 1.0000 | 0.8993 | | GPU utilization mean/max (%) | 93.9 | 43.3 | | Power mean/max (W) | 179.1 | 124.3 | Positive findings: peak inference memory is lower (generation 9.005 vs 9.912 GB, −9.2%), and the model draws less power because it is not compute-bound. Negative findings, stated plainly: - **Decode is 2.2× slower** (12.42 vs 27.33 tok/s) and **prefill is 2.8× slower** (TTFT 1431 vs 503 ms), despite the sparse path reading fewer keys. The current implementation loops over query blocks in Python and issues many small kernels, so launch overhead dominates the FLOPs saved. **The sparse attention does not yet pay off on this hardware.** - **Attention key accesses are still 89.9% of full attention** at this sequence length. The reason is structural: the local window already covers 3 blocks (192 tokens) and the sparse path reads up to `top_k + 1 = 9` blocks from a grid that only has 16 blocks at 1024 tokens, so the selected set is almost the whole grid. Sparsity only becomes a real saving once the sequence is long relative to `top_k × block_size` (i.e. well beyond 10k tokens). - **Language modeling perplexity is clearly worse than the base model** at every length tested (ppl 82.3 vs 39.4 at 512; 50.5 vs 21.8 at 2048). This is the honest cost of shrinking the dense local window from the full prefix to 192 tokens while the new long-range branches are still very weakly trained. - **Standard benchmarks are roughly neutral but not better**: arc_easy improves (+0.070 acc_norm), winogrande is unchanged, while hellaswag (−0.015) and piqa (−0.010) regress slightly. ### 4.4 Sparsity Two measurements are reported because they answer different questions and are **not** interchangeable: | Scenario | Average blocks selected per query | Key access ratio vs full attention | |---|---|---| | Chunked forward over a 512-token validation window | 0.38 | 0.0938 | | Prefill of 1024 tokens (steady state) | up to `top_k` | 0.8993 | The first number averages over *all* query blocks including the early ones, which have no completed blocks available to select and therefore read nothing through the sparse path. The second is the steady-state ratio for later queries, and it is the one that matters for efficiency — see the note in section 4.3: at these sequence lengths the sparse path is not yet saving meaningful work. --- ## 5. Known Limitations 1. **Trained for very little.** Only 5.0M tokens (≈0.25 epoch). The new branches have not converged; ppl is well above the base model and decode is slower (see 3.6 and 4.1). Continuing to 200M tokens or more is required before the SSA layers can genuinely take over long-range modeling. 2. **The shared branch does not appear to help and was actively suppressed by the optimizer.** Its learnable gate *decreased* over training (0.0100 → 0.0129 at step 100 → 0.0122 at step 610) instead of growing, meaning the optimizer found the single prefix-mean summary not worth injecting. This is the single most important thing to change next: replace it with **one compressed vector per block** (same `O(T·T/B)` cost, far more information retained). 3. **Sparse attention is not yet a net win on this hardware.** It reads fewer keys but runs 2.2× slower because the implementation loops over query blocks in Python and issues many small kernels. It needs kernel-level batching (or a fused implementation) before the sparsity can translate into speed. 4. **Sparsity only pays off at long sequences.** With `top_k=8` and `block_size=64`, the sparse path can read up to 9 blocks = 576 tokens; at 1024 tokens the grid only has 16 blocks, so the selected set covers most of the context and the local window already covers the rest. Real savings require sequences well beyond 10k tokens. 5. **Routing quality is not fully validated.** The distribution of selected top-k blocks should be checked for degeneration (e.g. always selecting the same blocks). The non-differentiable-routing bug that would have made this *impossible* to learn has been fixed (see section 2), but the learned policy has not been analyzed in detail. 6. **Block size B = 64 was not ablated.** Limited by memory and the Windows WDDM watchdog on this machine; B ∈ {32, 64, 128} should be swept on a larger GPU. 7. **Document-boundary packing.** The current packing strategy places multiple documents in one sequence (separated by EOS). --- ## 6. License **Free for personal use; a paid license is required for commercial use.** - ✅ Personal study, research, teaching, hobby projects: **free**, no application required - ✅ Academic research with public publication: **free** (please cite the source) - 💰 Internal company/studio use, paid API/SaaS, product integration, client deliverables: **commercial license required** **Commercial licensing contact: novaweb6868@outlook.com** Full terms: [LICENSE.custom.md](./LICENSE.custom.md). This model is an architectural retrofit of `Qwen/Qwen3-0.6B-Base` (Apache License 2.0). This license governs only the newly added portions and does not alter the upstream component's original license. --- ## 7. Citation ```bibtex @misc{baihu-v1-flash, title = {BaiHu-V1-Flash: An SSA (Sparse-attention + SubQ) Retrofit of Qwen3-0.6B-Base}, author = {NovaAI6868}, year = {2026}, url = {https://huggingface.co/NovaAI6868/BaiHu-V1-Flash} } ```