The flex_attention + paged KV + CUDA graphs inference engine was
Qwen3-specific in a handful of places, but the underlying engine
(PageTable, PagedKVCache, manual forward walker, decode graph capture,
double-copy LoRA rollout, FA4 capability guard) reads only attributes
that LlamaAttention / LlamaModel also expose. This change makes the
engine run on both Qwen3 and Llama-3.2-3B-Instruct.
Attention forward factory:
- make_flex_qwen3_attention_forward -> make_flex_attention_forward
- Guard the per-head QK RMSNorm call behind hasattr(self, "q_norm").
Qwen3 has it, Llama does not. The Qwen3 path is byte-equivalent to
before: RMSNorm on [B, S, H, D] (per-head) then transpose.
- patch_qwen3_model -> patch_model_attention_forwards.
Chat template selection:
- New --chat_template {auto,grpo,native}. auto picks GRPO for Qwen3
and the tokenizer's shipped template otherwise. grpo forces GRPO
(matches prior Qwen3 baselines). native forces the tokenizer's own
template (Llama-3.2-Instruct only produces coherent completions
with its shipped Instruct template).
Stats JSON:
- backend: "qwen3_flex" -> "flex"; adds "model_name" so multi-arch
runs land in a single schema.
README: one paragraph noting Llama-3.2 support + the --chat_template
native flag.
Measured on B200 (sm_100), n_prompts 64, max_new_tokens 512, 5 rounds,
--capture_cudagraph, double-copy LoRA rank 32:
Qwen3-4B-Base bf16 3975 tok/s 44.2 GB
Qwen3-4B-Base bf16 + LoRA 3656 tok/s 52.0 GB
Qwen3-4B-Base 4bit + LoRA 1734 tok/s 40.6 GB
Llama-3.2-3B-Inst bf16 4216 tok/s 34.7 GB
Llama-3.2-3B-Inst bf16+L 4205 tok/s 40.9 GB
Llama-3.2-3B-Inst 4bit+L 1892 tok/s 31.7 GB
--verify_no_drift passes on both arches (base bit-identical across 10
perturb+refresh cycles, inference hash deterministic).
Llama runs the Triton flex_attention backend instead of FA4:
flash-attn-4 b9's sm_100 kernel raises a NoneType in
handle_block_sparse_empty_tile_correction_sm100 on Llama-3.2's head
shapes. Qwen3 is unaffected. Pass --no-fa4_prefill on Llama; auto-FA4
still enables on Qwen3.
5bfefc2377