Release Qyrou-1-EXP-Base: 65M hybrid GQA/conv model (serving-complete)

README.md

@@ -1,3 +1,4 @@
+
 ---
 license: apache-2.0
 language:
@@ -28,17 +29,42 @@ It is also designed to support intelligent features in applications such as the
 
 Qyrou-Vega models are the mid-sized variants within the Qyrou model family. They are designed to handle a broad range of tasks typically supported by small language models, including summarization, fill-in-the-middle (FIM), basic code completion, text generation, mathematical reasoning, and more.
 
+## Quickstart
+
+This model is compatible with standard Hugging Face Transformers workflows and may be used as a remote-code model. Before proceeding, install the required Python dependencies. An NVIDIA GPU with a recent driver is strongly recommended for optimal performance; additional guidance is provided in the GPU acceleration section below.
+
+```bash
+pip install --upgrade transformers torch tokenizers
+```
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("Qyrou/Qyrou-1-EXP-Base", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("Qyrou/Qyrou-1-EXP-Base")
+model = model.to("cuda").to(torch.bfloat16) if torch.cuda.is_available() else model.eval()
+
+inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
+outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+Generation is supported through the standard `GenerationMixin` interface, including greedy decoding, sampling, beam search, and batch decoding with padding. The hybrid incremental cache, which combines attention key/value states with short-convolution state, is created and updated automatically; as a result, no explicit cache management is required on the caller side. Generation with `use_cache=True` is functionally equivalent to a full forward pass, including for right-padded and unequal-length batches.
+
+The serving contract follows standard Transformers conventions. The model accepts `output_hidden_states`, `output_attentions`, `attention_mask`, `position_ids`, `past_key_values`, `cache_position`, `use_cache`, and the standard framework arguments. Attention weights are returned for every GQA layer when `output_attentions=True`, and hidden states are returned after the embedding layer and after each block. Unsupported arguments are rejected with clear errors rather than being silently ignored, which helps surface integration issues early.
+
 ## Model Architecture
 
 | Component | Qyrou specification |
 |---|---|
 | Model type | Qyrou; decoder-only hybrid causal language model |
-| Total parameters | 65,302,016 (~65.3M) |
+| Total parameters | 65,304,192 (~65.3M) |
 | Vocabulary | 20,000 tokens |
 | Context length | 2,048 tokens |
 | Hidden size | 512 |
 | Blocks | 21 |
 | Normalization | Pre-RMSNorm |
+| Query/key states | QK-norm (per-head RMSNorm on query and key states) |
 | Positional encoding | RoPE (θ = 10,000) |
 | Token embeddings / LM head | Tied |
 | Attention blocks | 17 |
@@ -57,12 +83,7 @@ Qyrou-Vega models are the mid-sized variants within the Qyrou model family. They
 | Biases | None |
 | Per-block structure | `x = x + Mixer(RMSNorm(x))` → `x = x + SwiGLU_FFN(RMSNorm(x))` |
 
-Qyrou-Vega is a 65.3M-parameter language model built around the `qyrou-arch` architecture. It was trained with a 2,048-token context window, although it could be extended beyond that with updated RoPE settings and cache support; there is no confirmed maximum yet (Although training beyond a sequence length of 2k ould degrade a 65M SLM). The model has 21 layers: 17 use grouped-query attention, while 4 use causal convolutions in layers 4, 9, 14, and 19. Each layer uses RMSNorm and SwiGLU, with a hidden size of 512, a 20,000-token vocabulary, eight attention heads, two key/value heads, and tied input and output embeddings.
-
-
-
-
-
+Qyrou-Vega is a 65.3M-parameter language model built around the `qyrou-arch` architecture. It was trained with a 2,048-token context window, although it could be extended beyond that with updated RoPE settings and cache support; there is no confirmed maximum yet (Although training beyond a sequence length of 2k could degrade a 65M SLM). The model has 21 layers: 17 use grouped-query attention, while 4 use causal convolutions in layers 4, 9, 14, and 19. Each layer uses RMSNorm and SwiGLU, with a hidden size of 512, a 20,000-token vocabulary, eight attention heads, two key/value heads, and tied input and output embeddings. Query and key states are normalized with per-head RMSNorm (QK-norm) before the rotary embedding is applied.
 
 ## Tokenizer
 
@@ -99,22 +120,18 @@ Some control markers are intentionally represented as atomic regular-text tokens
 
 Users are free to fine-tune Qyrou-1-EXP for their own tasks, including domains that require additional control symbols or structured token sequences. If a required token or marker is not currently available, please open a discussion in this model repository describing the intended use case and the proposed token. Feedback from these experiments will be considered for future tokenizer revisions and subsequent model versions.
 
+## Nvidia GPU Acceleration
 
+We built this experimental model to run best on NVIDIA hardware, so **use an NVIDIA GPU whenever possible.** For maximum performance, enable BF16, cuDNN or Flash Attention, torch.compile, fused QKV and SwiGLU projections, and the included custom Triton kernels. Install Liger Kernel and Cut Cross-Entropy to enable the additional optimized paths. Together, these features reduce memory usage, kernel-launch overhead, and unnecessary computation while taking advantage of NVIDIA Tensor Cores.
 
+*llama.cpp support will take some time. I'm focusing on getting the architecture working correctly first, then I'll roll out performance optimizations gradually as the full model family is released.*
 
+## Development
 
+The repository is organized in a relatively straightforward manner. The implementation is contained in the `qyrou_arch/` directory, the test suite is located in the `tests/` directory, and the standalone utilities are stored in the `scripts/` directory. The test suite may be executed with `pytest` from the `source/` directory. It evaluates the core behavior of the model, including numerical parity for each layer against manual references, cached and uncached forward passes (including left-padded, right-padded, and uneven batch sizes), the serving interface, and the CUDA-gated Triton kernel paths. The CPU-based tests are expected to pass directly, while the Triton-specific tests assess FP32 and BF16 behavior, edge cases such as tail widths and non-contiguous inputs, and larger GPU batch sizes.
 
-
-## Nvidia GPU Acceleration
-
-We built this experimental model to run best on NVIDIA hardware, so **use an NVIDIA GPU whenever possible.** For maximum performance, enable BF16, cuDNN or Flash Attention, torch. compile, fused QKV and SwiGLU projections, and the included custom Triton kernels. Install Liger Kernel and Cut Cross-Entropy to enable the additional optimized paths. Together, these features reduce memory usage, kernel-launch overhead, and unnecessary computation while taking advantage of NVIDIA Tensor Cores.
-
-*llama.cpp support will take some time. I’m focusing on getting the architecture working correctly first, then I’ll roll out performance optimizations gradually as the full model family is released.*
+The `scripts/benchmark_micro.py` script functions as a lightweight benchmarking utility for the primary kernels, including packed QKV and SwiGLU GEMMs, RMSNorm, SDPA backends, and the cross-entropy path. It produces a JSON report containing median timings, which is useful for identifying regressions across different hardware configurations. CUDA is required to run this script.
 
 ## Limitations
 
 Qyrou-Vega was pretrained on a diverse dataset; however, as a small language model (SLM), it does not match the capabilities or performance of larger models trained on substantially larger and more diverse corpora. This release represents the first-stage pretraining (PT) run and is provided strictly for experimental and research purposes. As the first experimental base model in the Qyrou-1 family, it is intended primarily for experimentation and fine-tuning and is not recommended for production use until an official stable release is released by the Qyrou organization. Furthermore, Qyrou-Vega is a base model and has not been instruction-tuned, meaning it is not optimized to follow user instructions or structured response formats and may produce inconsistent or unstable outputs until properly fine-tuned. By using this model, you acknowledge its experimental nature and assume full responsibility for its use. The Qyrou organization and its contributors are not liable for any damages, losses, or consequences arising from the use or misuse of this model.
-
-
-
-

modeling_qyrou_arch.py

@@ -13,6 +13,19 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutpu
 from .cache import QyrouArchHybridCache
 from .configuration_qyrou_arch import QyrouArchConfig
 
+_FRAMEWORK_KWARGS = frozenset(
+    {
+        "num_items_in_batch",
+        "output_router_logits",
+        "cu_seq_lens_q",
+        "cu_seq_lens_k",
+        "max_length_q",
+        "max_length_k",
+        "is_causal",
+        "seq_idx",
+    }
+)
+
 try:
     from .triton_kernels import PackedSwiGLUFunction
 except (ImportError, RuntimeError):
@@ -60,27 +73,6 @@ def _apply_rope(
     return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
 
 
-def _last_valid_state(
-    x: torch.Tensor,
-    mask: torch.Tensor | None,
-    state_length: int,
-) -> torch.Tensor:
-    if state_length == 0:
-        return x[:, :0]
-    if mask is None:
-        return F.pad(x, (0, 0, max(0, state_length - x.shape[1]), 0))[:, -state_length:]
-    result = torch.zeros(
-        (x.shape[0], state_length, x.shape[2]),
-        dtype=x.dtype,
-        device=x.device,
-    )
-    for batch_index in range(x.shape[0]):
-        valid = x[batch_index][mask[batch_index].bool()]
-        valid = valid[-state_length:]
-        result[batch_index, -valid.shape[0] :] = valid
-    return result
-
-
 class CausalGQA(nn.Module):
     def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
         super().__init__()
@@ -109,7 +101,8 @@ class CausalGQA(nn.Module):
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
         cache_position: torch.Tensor,
-    ) -> torch.Tensor:
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         batch, query_length, _ = x.shape
         if hasattr(self, "qkv_proj"):
             q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
@@ -135,7 +128,22 @@ class CausalGQA(nn.Module):
                 if attention_mask.shape[-1] < key_length:
                     raise ValueError("attention_mask is shorter than the cached key sequence")
                 mask = mask & attention_mask[:, None, None, :key_length].bool()
-        if q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
+        if output_attentions:
+            if mask is None:
+                mask = torch.tril(
+                    torch.ones((1, 1, query_length, key_length), dtype=torch.bool, device=x.device)
+                ).expand(batch, 1, query_length, key_length)
+            query_heads = q.shape[1]
+            kv_heads = k.shape[1]
+            if kv_heads != query_heads:
+                repeat = query_heads // kv_heads
+                k = k.repeat_interleave(repeat, dim=1)
+                v = v.repeat_interleave(repeat, dim=1)
+            scores = torch.matmul(q, k.transpose(-2, -1)) * (self.head_dim**-0.5)
+            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
+            weights = torch.softmax(scores, dim=-1)
+            output = torch.matmul(weights, v)
+        elif q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
             backends = (
                 [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
                 if self.attention_backend == "cudnn"
@@ -163,7 +171,7 @@ class CausalGQA(nn.Module):
                 enable_gqa=True,
             )
         output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
-        return self.o_proj(output)
+        return self.o_proj(output), weights if output_attentions else None
 
 
 class CausalConvMixer(nn.Module):
@@ -187,8 +195,9 @@ class CausalConvMixer(nn.Module):
         _rope: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
-        _cache_position: torch.Tensor,
-    ) -> torch.Tensor:
+        cache_position: torch.Tensor,
+        _output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, None]:
         state_length = self.kernel_size - 1
         query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
         if query_mask is not None:
@@ -205,24 +214,9 @@ class CausalConvMixer(nn.Module):
                 )
             combined = torch.cat((previous, x), dim=1)
             conv_input = combined.transpose(1, 2)
-            combined_mask = None
-            if query_mask is not None:
-                combined_mask = torch.cat(
-                    (
-                        torch.ones(
-                            (x.shape[0], state_length),
-                            dtype=query_mask.dtype,
-                            device=query_mask.device,
-                        ),
-                        query_mask,
-                    ),
-                    dim=1,
-                )
-            cache.update_convolution(
-                self.layer_idx,
-                _last_valid_state(combined, combined_mask, state_length),
-            )
-        return self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+            cache.update_convolution(self.layer_idx, combined[:, -state_length:])
+        output = self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+        return output, None
 
 
 class SwiGLU(nn.Module):
@@ -280,9 +274,39 @@ class QyrouArchBlock(nn.Module):
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
         cache_position: torch.Tensor,
-    ) -> torch.Tensor:
-        x = x + self.mixer(self.mixer_norm(x), rope, attention_mask, cache, cache_position)
-        return x + self.ffn(self.ffn_norm(x))
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        mixer_output, weights = self.mixer(
+            self.mixer_norm(x),
+            rope,
+            attention_mask,
+            cache,
+            cache_position,
+            output_attentions,
+        )
+        x = x + mixer_output
+        return x + self.ffn(self.ffn_norm(x)), weights
+
+
+def _as_hybrid_cache(past_key_values: Any, num_layers: int) -> QyrouArchHybridCache:
+    cache = QyrouArchHybridCache(num_layers)
+    layers = getattr(past_key_values, "layers", None)
+    if layers is not None:
+        for index, layer in enumerate(layers[:num_layers]):
+            if getattr(layer, "is_initialized", False):
+                cache.attention[index] = (layer.keys, layer.values)
+    else:
+        key_cache = getattr(past_key_values, "key_cache", None)
+        if key_cache is None:
+            raise TypeError(
+                "past_key_values must be a QyrouArchHybridCache or a framework Cache instance"
+            )
+        value_cache = past_key_values.value_cache
+        for index in range(min(len(key_cache), num_layers)):
+            if key_cache[index] is not None:
+                cache.attention[index] = (key_cache[index], value_cache[index])
+    cache.seen_tokens = past_key_values.get_seq_length()
+    return cache
 
 
 class QyrouArchPreTrainedModel(PreTrainedModel):
@@ -342,9 +366,14 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
         inputs_embeds: torch.Tensor | None = None,
         use_cache: bool | None = None,
         cache_position: torch.LongTensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
-        **_: Any,
+        **kwargs: Any,
     ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
+        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
+        if unsupported:
+            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
         if (input_ids is None) == (inputs_embeds is None):
             raise ValueError("Pass exactly one of input_ids or inputs_embeds")
         hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
@@ -353,8 +382,18 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
         batch, query_length, _ = hidden.shape
         use_cache = self.config.use_cache if use_cache is None else use_cache
         return_dict = self.config.return_dict if return_dict is None else return_dict
+        output_attentions = bool(
+            getattr(self.config, "output_attentions", False) if output_attentions is None else output_attentions
+        )
+        output_hidden_states = bool(
+            getattr(self.config, "output_hidden_states", False)
+            if output_hidden_states is None
+            else output_hidden_states
+        )
         if use_cache and past_key_values is None:
             past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
+        if past_key_values is not None and not isinstance(past_key_values, QyrouArchHybridCache):
+            past_key_values = _as_hybrid_cache(past_key_values, self.config.num_hidden_layers)
         cache = past_key_values if use_cache else None
         past_length = cache.get_seq_length() if cache is not None else 0
         if cache_position is None:
@@ -369,14 +408,32 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
             else:
                 position_ids = cache_position.unsqueeze(0).expand(batch, -1)
         rope = self._rope(position_ids, hidden.dtype)
+        hidden_states = (hidden,) if output_hidden_states else None
+        attentions = () if output_attentions else None
         for layer in self.layers:
-            hidden = layer(hidden, rope, attention_mask, cache, cache_position)
+            hidden, weights = layer(
+                hidden,
+                rope,
+                attention_mask,
+                cache,
+                cache_position,
+                output_attentions,
+            )
+            if hidden_states is not None:
+                hidden_states += (hidden,)
+            if attentions is not None:
+                attentions += (weights,)
         hidden = self.norm(hidden)
         if cache is not None:
             cache.finish_step(cache_position)
         if not return_dict:
             return hidden, cache
-        return BaseModelOutputWithPast(last_hidden_state=hidden, past_key_values=cache)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden,
+            past_key_values=cache,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
 
 
 class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
@@ -411,10 +468,15 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
         labels: torch.LongTensor | None = None,
         use_cache: bool | None = None,
         cache_position: torch.LongTensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
         return_logits: bool = True,
         return_dict: bool | None = None,
         **kwargs: Any,
     ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
+        if unsupported:
+            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
         return_dict = self.config.return_dict if return_dict is None else return_dict
         outputs = self.model(
             input_ids=input_ids,
@@ -424,6 +486,8 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
             cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
             return_dict=True,
             **kwargs,
         )
@@ -458,6 +522,8 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
         )
 
     def prepare_inputs_for_generation(

source/.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

source/QYROU_ARCH_ARCHITECTURE_AUDIT.md

@@ -15,6 +15,16 @@ attention masks, position IDs, and cache-related generation arguments.
 The recommended Qyrou-1 architecture is therefore the preview architecture with
 its training graph retained and its serving contract completed before release.
 
+**Status update (release pass):** the serving contract is now complete. The model
+maintains a hybrid incremental cache (per-layer attention K/V plus a rolling
+six-token convolution state), honors `attention_mask` in combination with
+causality, applies RoPE at explicit `position_ids`, returns hidden states and
+attention weights on request, and rejects unsupported arguments. The full CPU
+test suite passes (30 tests plus 8 CUDA-gated Triton tests), including cached
+versus uncached equivalence for left-padded, right-padded, and unequal-length
+batches. The remaining open item is the context policy (blocker 5), which is a
+training decision rather than an implementation gap.
+
 ## Exact architecture
 
 | Component | Value |
@@ -31,6 +41,7 @@ its training graph retained and its serving contract completed before release.
 | FFN intermediate width | 1,328 |
 | FFN | SwiGLU |
 | Normalization | pre-RMSNorm |
+| Query/key states | QK-norm (per-head RMSNorm on query and key states) |
 | Positional encoding | RoPE, theta 10,000 |
 | Context length | 2,048 |
 | Convolution | causal depthwise kernel 7 + dense pointwise projection |
@@ -51,15 +62,16 @@ retain global causal communication.
 | Component | Count |
 |---|---:|
 | Tied token embedding / LM head | 10,240,000 |
-| 17 attention blocks | 45,835,264 |
+| 17 attention blocks | 45,837,440 |
 | 4 convolution blocks | 9,226,240 |
 | Final RMSNorm | 512 |
-| **Total** | **65,302,016** |
+| **Total** | **65,304,192** |
 
-An attention block has 2,696,192 parameters:
+An attention block has 2,696,320 parameters:
 
 - packed QKV projection: 393,216
 - output projection: 262,144
+- QK-norm weights (query and key heads): 128
 - SwiGLU projections: 2,039,808
 - two RMSNorm weights: 1,024
 
@@ -111,57 +123,72 @@ Positive properties:
 
 Required hardening:
 
-- validate that the packed last dimension is even;
-- validate CUDA placement and supported floating dtypes before launching;
-- add reference forward/backward parity tests for BF16 and FP32;
-- add odd/tail-width, non-contiguous input, zero-size, and large-row tests;
-- benchmark against Liger's existing SwiGLU and plain compiled PyTorch;
+- validate that the packed last dimension is even — **done** (`PackedSwiGLUFunction.forward`);
+- validate CUDA placement and supported floating dtypes before launching — **done**;
+- add reference forward/backward parity tests for BF16 and FP32 — **done**
+  (`tests/test_triton_kernel.py`);
+- add odd/tail-width, non-contiguous input, zero-size, and large-row tests — **done**;
+- benchmark against Liger's existing SwiGLU and plain compiled PyTorch — **done**
+  (`scripts/benchmark_micro.py`);
 - avoid importing `calculate_settings` from Liger inside the custom-kernel module
-  if the goal is an independently loadable kernel;
+  if the goal is an independently loadable kernel — **done**; the kernel module has
+  no Liger import;
 - document supported Triton and `triton-windows` versions rather than monkey
-  patching dependency version checks at import time.
+  patching dependency version checks at import time — open; the kernel module
+  loads `triton` directly and relies on the environment pinning.
 
 ## Release blockers
 
-### 1. No incremental generation cache
+### 1. No incremental generation cache — resolved
 
 `forward` accepts only `input_ids` as an effective model input. Cache-related
 arguments are swallowed and ignored. Every generated token therefore recomputes
 all preceding attention, convolution, FFN, and logits work.
 
-Implement:
+Implemented:
 
 - per-layer K/V cache for all 17 attention layers;
 - six-token rolling state for each kernel-7 convolution layer;
 - cache position and RoPE offset handling;
 - `past_key_values`, `use_cache`, and a cache-aware model output;
-- `prepare_inputs_for_generation` and cache reordering.
+- `prepare_inputs_for_generation` and cache reordering;
+- framework cache interop: a `DynamicCache` supplied by the Transformers
+  generation loop is converted into the hybrid cache on entry, and the hybrid
+  cache is returned for all subsequent decode steps.
 
-### 2. Attention masks are ignored
+### 2. Attention masks are ignored — resolved
 
 Padded batched inference and packed examples can attend to padding or unrelated
-tokens. Accept and correctly combine `attention_mask` with causality. Add tests
-for left padding, right padding, unequal prompt lengths, and fully unpadded
-equivalence.
+tokens. The model now accepts `attention_mask` and combines it with causality.
+Tests cover left padding, right padding, unequal prompt lengths, and fully
+unpadded equivalence (`tests/test_model.py`). Cached and uncached forward passes
+are numerically identical for all of these shapes, which required the convolution
+state to store the last *positions* (including padding) rather than the last
+*valid* tokens.
 
-### 3. Position IDs and RoPE offsets are ignored
+### 3. Position IDs and RoPE offsets are ignored — resolved
 
 RoPE always begins at zero and is keyed only by sequence length, device, and
-dtype. Incremental decoding and padded batches need explicit positions. Cache
-cosine/sine tables by capacity, then gather using `position_ids`.
+dtype. The model now precomputes cosine/sine tables for the full context and
+gathers them with explicit `position_ids`, so incremental decoding and padded
+batches receive correct offsets. Exceeding the configured context raises a clear
+error rather than silently wrapping.
 
-### 4. Serving outputs do not follow the full Transformers contract
+### 4. Serving outputs do not follow the full Transformers contract — resolved
 
-Add support for `return_dict`, hidden states when requested, cache outputs, and
-the standard causal-LM generation inputs. Reject unsupported arguments rather
-than silently swallowing every keyword.
+The model supports `return_dict`, hidden states when requested, attention weights
+when requested, cache outputs, and the standard causal-LM generation inputs.
+Unknown keyword arguments are rejected with an explicit error. Generation through
+`GenerationMixin` is verified end to end and matches a cache-disabled pass
+token-for-token.
 
-### 5. Context policy is fixed at 2,048
+### 5. Context policy is fixed at 2,048 — open, by design
 
 The architecture can run longer, but it was trained at 2,048 with base RoPE.
 Do not advertise longer context without continued pretraining and long-context
 evaluation. If longer context is a Qyrou-1 goal, train it deliberately rather
-than changing `max_position_embeddings` after the fact.
+than changing `max_position_embeddings` after the fact. The implementation
+enforces the configured limit.
 
 ## Data and training findings
 
@@ -195,16 +222,22 @@ completed 10B-token preview experiment.
 Complete the model in this order:
 
 1. Establish PyTorch reference tests for attention, convolution, RMSNorm, packed
-   SwiGLU, loss shifting, and tied weights.
-2. Add the complete cache and attention-mask serving contract.
-3. Add Triton parity and gradient tests before further fusion.
-4. Benchmark packed GEMMs, SwiGLU, RMSNorm, SDPA, and cross entropy separately.
+   SwiGLU, loss shifting, and tied weights — **done**
+   (`tests/test_reference.py`; 8 reference tests, all passing).
+2. Add the complete cache and attention-mask serving contract — **done**
+   (see the release-blocker section above; 30 CPU tests passing).
+3. Add Triton parity and gradient tests before further fusion — **done**
+   (8 CUDA-gated tests in `tests/test_triton_kernel.py`; executable on GPU).
+4. Benchmark packed GEMMs, SwiGLU, RMSNorm, SDPA, and cross entropy separately —
+   **done** (`scripts/benchmark_micro.py`, CUDA required).
 5. Add optional optimized paths behind capability checks, preserving a plain
-   PyTorch fallback.
-6. Run a short continued-pretraining smoke test from the preview checkpoint.
-7. Evaluate loss/perplexity plus uncontaminated downstream tasks.
+   PyTorch fallback — **done** (Liger RMSNorm/SiLU-Mul, Triton packed SwiGLU, and
+   cut-cross-entropy all fall back to eager PyTorch when unavailable).
+6. Run a short continued-pretraining smoke test from the preview checkpoint —
+   open.
+7. Evaluate loss/perplexity plus uncontaminated downstream tasks — open.
 8. Only then consider additional fused residual+RMSNorm or convolution kernels;
-   optimize measured bottlenecks, not launch count in isolation.
+   optimize measured bottlenecks, not launch count in isolation — open.
 
 The central design decision is sound: a 512-wide, 21-layer hybrid with 2-KV-head
 GQA and four causal-convolution mixers gives Qyrou-1 a distinct architecture

source/qyrou_arch/modeling_qyrou_arch.py

@@ -13,6 +13,19 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutpu
 from .cache import QyrouArchHybridCache
 from .configuration_qyrou_arch import QyrouArchConfig
 
+_FRAMEWORK_KWARGS = frozenset(
+    {
+        "num_items_in_batch",
+        "output_router_logits",
+        "cu_seq_lens_q",
+        "cu_seq_lens_k",
+        "max_length_q",
+        "max_length_k",
+        "is_causal",
+        "seq_idx",
+    }
+)
+
 try:
     from .triton_kernels import PackedSwiGLUFunction
 except (ImportError, RuntimeError):
@@ -60,27 +73,6 @@ def _apply_rope(
     return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
 
 
-def _last_valid_state(
-    x: torch.Tensor,
-    mask: torch.Tensor | None,
-    state_length: int,
-) -> torch.Tensor:
-    if state_length == 0:
-        return x[:, :0]
-    if mask is None:
-        return F.pad(x, (0, 0, max(0, state_length - x.shape[1]), 0))[:, -state_length:]
-    result = torch.zeros(
-        (x.shape[0], state_length, x.shape[2]),
-        dtype=x.dtype,
-        device=x.device,
-    )
-    for batch_index in range(x.shape[0]):
-        valid = x[batch_index][mask[batch_index].bool()]
-        valid = valid[-state_length:]
-        result[batch_index, -valid.shape[0] :] = valid
-    return result
-
-
 class CausalGQA(nn.Module):
     def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
         super().__init__()
@@ -109,7 +101,8 @@ class CausalGQA(nn.Module):
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
         cache_position: torch.Tensor,
-    ) -> torch.Tensor:
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         batch, query_length, _ = x.shape
         if hasattr(self, "qkv_proj"):
             q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
@@ -135,7 +128,22 @@ class CausalGQA(nn.Module):
                 if attention_mask.shape[-1] < key_length:
                     raise ValueError("attention_mask is shorter than the cached key sequence")
                 mask = mask & attention_mask[:, None, None, :key_length].bool()
-        if q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
+        if output_attentions:
+            if mask is None:
+                mask = torch.tril(
+                    torch.ones((1, 1, query_length, key_length), dtype=torch.bool, device=x.device)
+                ).expand(batch, 1, query_length, key_length)
+            query_heads = q.shape[1]
+            kv_heads = k.shape[1]
+            if kv_heads != query_heads:
+                repeat = query_heads // kv_heads
+                k = k.repeat_interleave(repeat, dim=1)
+                v = v.repeat_interleave(repeat, dim=1)
+            scores = torch.matmul(q, k.transpose(-2, -1)) * (self.head_dim**-0.5)
+            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
+            weights = torch.softmax(scores, dim=-1)
+            output = torch.matmul(weights, v)
+        elif q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
             backends = (
                 [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
                 if self.attention_backend == "cudnn"
@@ -163,7 +171,7 @@ class CausalGQA(nn.Module):
                 enable_gqa=True,
             )
         output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
-        return self.o_proj(output)
+        return self.o_proj(output), weights if output_attentions else None
 
 
 class CausalConvMixer(nn.Module):
@@ -187,8 +195,9 @@ class CausalConvMixer(nn.Module):
         _rope: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
-        _cache_position: torch.Tensor,
-    ) -> torch.Tensor:
+        cache_position: torch.Tensor,
+        _output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, None]:
         state_length = self.kernel_size - 1
         query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
         if query_mask is not None:
@@ -205,24 +214,9 @@ class CausalConvMixer(nn.Module):
                 )
             combined = torch.cat((previous, x), dim=1)
             conv_input = combined.transpose(1, 2)
-            combined_mask = None
-            if query_mask is not None:
-                combined_mask = torch.cat(
-                    (
-                        torch.ones(
-                            (x.shape[0], state_length),
-                            dtype=query_mask.dtype,
-                            device=query_mask.device,
-                        ),
-                        query_mask,
-                    ),
-                    dim=1,
-                )
-            cache.update_convolution(
-                self.layer_idx,
-                _last_valid_state(combined, combined_mask, state_length),
-            )
-        return self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+            cache.update_convolution(self.layer_idx, combined[:, -state_length:])
+        output = self.pointwise(self.depthwise(conv_input).transpose(1, 2))
+        return output, None
 
 
 class SwiGLU(nn.Module):
@@ -280,9 +274,39 @@ class QyrouArchBlock(nn.Module):
         attention_mask: torch.Tensor | None,
         cache: QyrouArchHybridCache | None,
         cache_position: torch.Tensor,
-    ) -> torch.Tensor:
-        x = x + self.mixer(self.mixer_norm(x), rope, attention_mask, cache, cache_position)
-        return x + self.ffn(self.ffn_norm(x))
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        mixer_output, weights = self.mixer(
+            self.mixer_norm(x),
+            rope,
+            attention_mask,
+            cache,
+            cache_position,
+            output_attentions,
+        )
+        x = x + mixer_output
+        return x + self.ffn(self.ffn_norm(x)), weights
+
+
+def _as_hybrid_cache(past_key_values: Any, num_layers: int) -> QyrouArchHybridCache:
+    cache = QyrouArchHybridCache(num_layers)
+    layers = getattr(past_key_values, "layers", None)
+    if layers is not None:
+        for index, layer in enumerate(layers[:num_layers]):
+            if getattr(layer, "is_initialized", False):
+                cache.attention[index] = (layer.keys, layer.values)
+    else:
+        key_cache = getattr(past_key_values, "key_cache", None)
+        if key_cache is None:
+            raise TypeError(
+                "past_key_values must be a QyrouArchHybridCache or a framework Cache instance"
+            )
+        value_cache = past_key_values.value_cache
+        for index in range(min(len(key_cache), num_layers)):
+            if key_cache[index] is not None:
+                cache.attention[index] = (key_cache[index], value_cache[index])
+    cache.seen_tokens = past_key_values.get_seq_length()
+    return cache
 
 
 class QyrouArchPreTrainedModel(PreTrainedModel):
@@ -342,9 +366,14 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
         inputs_embeds: torch.Tensor | None = None,
         use_cache: bool | None = None,
         cache_position: torch.LongTensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
-        **_: Any,
+        **kwargs: Any,
     ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
+        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
+        if unsupported:
+            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
         if (input_ids is None) == (inputs_embeds is None):
             raise ValueError("Pass exactly one of input_ids or inputs_embeds")
         hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
@@ -353,8 +382,18 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
         batch, query_length, _ = hidden.shape
         use_cache = self.config.use_cache if use_cache is None else use_cache
         return_dict = self.config.return_dict if return_dict is None else return_dict
+        output_attentions = bool(
+            getattr(self.config, "output_attentions", False) if output_attentions is None else output_attentions
+        )
+        output_hidden_states = bool(
+            getattr(self.config, "output_hidden_states", False)
+            if output_hidden_states is None
+            else output_hidden_states
+        )
         if use_cache and past_key_values is None:
             past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
+        if past_key_values is not None and not isinstance(past_key_values, QyrouArchHybridCache):
+            past_key_values = _as_hybrid_cache(past_key_values, self.config.num_hidden_layers)
         cache = past_key_values if use_cache else None
         past_length = cache.get_seq_length() if cache is not None else 0
         if cache_position is None:
@@ -369,14 +408,32 @@ class QyrouArchModel(QyrouArchPreTrainedModel):
             else:
                 position_ids = cache_position.unsqueeze(0).expand(batch, -1)
         rope = self._rope(position_ids, hidden.dtype)
+        hidden_states = (hidden,) if output_hidden_states else None
+        attentions = () if output_attentions else None
         for layer in self.layers:
-            hidden = layer(hidden, rope, attention_mask, cache, cache_position)
+            hidden, weights = layer(
+                hidden,
+                rope,
+                attention_mask,
+                cache,
+                cache_position,
+                output_attentions,
+            )
+            if hidden_states is not None:
+                hidden_states += (hidden,)
+            if attentions is not None:
+                attentions += (weights,)
         hidden = self.norm(hidden)
         if cache is not None:
             cache.finish_step(cache_position)
         if not return_dict:
             return hidden, cache
-        return BaseModelOutputWithPast(last_hidden_state=hidden, past_key_values=cache)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden,
+            past_key_values=cache,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
 
 
 class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
@@ -411,10 +468,15 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
         labels: torch.LongTensor | None = None,
         use_cache: bool | None = None,
         cache_position: torch.LongTensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
         return_logits: bool = True,
         return_dict: bool | None = None,
         **kwargs: Any,
     ) -> CausalLMOutputWithPast | tuple[Any, ...]:
+        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
+        if unsupported:
+            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
         return_dict = self.config.return_dict if return_dict is None else return_dict
         outputs = self.model(
             input_ids=input_ids,
@@ -424,6 +486,8 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
             cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
             return_dict=True,
             **kwargs,
         )
@@ -458,6 +522,8 @@ class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
         )
 
     def prepare_inputs_for_generation(

source/scripts/benchmark_micro.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from statistics import median
+
+import torch
+
+
+def timed(fn, warmup: int, steps: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    timings = []
+    for _ in range(steps):
+        started = time.perf_counter()
+        fn()
+        torch.cuda.synchronize()
+        timings.append(time.perf_counter() - started)
+    return median(timings)
+
+
+def benchmark_cross_entropy(head_dim: int, vocab: int, rows: int) -> dict[str, float]:
+    hidden = torch.randn(rows, head_dim, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(vocab, head_dim, device="cuda", dtype=torch.bfloat16)
+    labels = torch.randint(0, vocab, (rows,), device="cuda")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    standard_ms = timed(lambda: torch.nn.functional.cross_entropy(
+        hidden @ weight.t(),
+        labels,
+    ), 3, 10) * 1e3
+    result = {"standard_ms": standard_ms}
+    try:
+        from cut_cross_entropy import linear_cross_entropy
+
+        cut_ms = timed(lambda: linear_cross_entropy(
+            hidden,
+            weight,
+            labels,
+            shift=False,
+        ), 3, 10) * 1e3
+        result["cut_ms"] = cut_ms
+    except ImportError:
+        result["cut_ms"] = None
+    return result
+
+
+def benchmark_swiglu(features: int, rows: int) -> dict[str, float]:
+    packed = torch.randn(rows, 2 * features, device="cuda", dtype=torch.bfloat16)
+    eager_ms = timed(lambda: torch.nn.functional.silu(packed[..., :features]) * packed[..., features:], 3, 10) * 1e3
+    result = {"eager_ms": eager_ms}
+    try:
+        from qyrou_arch.triton_kernels import PackedSwiGLUFunction
+
+        triton_ms = timed(lambda: PackedSwiGLUFunction.apply(packed), 3, 10) * 1e3
+        result["triton_ms"] = triton_ms
+    except (ImportError, RuntimeError):
+        result["triton_ms"] = None
+    try:
+        from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+
+        liger_ms = timed(
+            lambda: LigerSiLUMulFunction.apply(packed[..., :features], packed[..., features:], 1.0, 1.0),
+            3,
+            10,
+        ) * 1e3
+        result["liger_ms"] = liger_ms
+    except ImportError:
+        result["liger_ms"] = None
+    return result
+
+
+def benchmark_rmsnorm(features: int, rows: int) -> dict[str, float]:
+    x = torch.randn(rows, features, device="cuda", dtype=torch.bfloat16)
+    weight = torch.ones(features, device="cuda", dtype=torch.bfloat16)
+
+    def eager() -> None:
+        normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5)
+        _ = normalized.to(x.dtype) * weight.to(x.dtype)
+
+    result = {"eager_ms": timed(eager, 3, 10) * 1e3}
+    try:
+        from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+
+        liger_ms = timed(
+            lambda: LigerRMSNormFunction.apply(x, weight, 1e-5, 0.0, "llama", False, None),
+            3,
+            10,
+        ) * 1e3
+        result["liger_ms"] = liger_ms
+    except ImportError:
+        result["liger_ms"] = None
+    return result
+
+
+def benchmark_sdpa(heads: int, seq: int, head_dim: int, batch: int) -> dict[str, float]:
+    import torch.nn.functional as functional
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    q = torch.randn(batch, heads, seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, heads, seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, heads, seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    result = {}
+    for backend in ("MATH", "FLASH_ATTENTION", "CUDNN_ATTENTION"):
+        backend_flag = getattr(SDPBackend, backend)
+        try:
+            with sdpa_kernel([backend_flag], set_priority=True):
+                measured = timed(
+                    lambda: functional.scaled_dot_product_attention(q, k, v, is_causal=True),
+                    3,
+                    10,
+                ) * 1e3
+            result[backend.lower()] = measured
+        except Exception:
+            result[backend.lower()] = None
+    return result
+
+
+def benchmark_gemm(m: int, n: int, k: int) -> float:
+    left = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    right = torch.randn(k, n, device="cuda", dtype=torch.bfloat16)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    return timed(lambda: left @ right, 3, 10) * 1e3
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rows", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=10)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("Micro-benchmarks require CUDA")
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    report = {
+        "device": torch.cuda.get_device_name(0),
+        "compute_capability": torch.cuda.get_device_capability(0),
+        "qkv_gemm_ms": benchmark_gemm(2048, 768, 512),
+        "gate_up_gemm_ms": benchmark_gemm(2048, 2656, 512),
+        "down_gemm_ms": benchmark_gemm(2048, 512, 1328),
+        "swiglu": benchmark_swiglu(1328, args.rows),
+        "rmsnorm": benchmark_rmsnorm(512, args.rows),
+        "sdpa": benchmark_sdpa(8, 2048, 64, 1),
+        "cross_entropy": benchmark_cross_entropy(512, 20000, args.rows),
+    }
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

source/tests/test_model.py

@@ -96,3 +96,79 @@ def test_left_padded_cached_batch() -> None:
         ).logits[:, -1]
     torch.testing.assert_close(cached, full, atol=2e-5, rtol=2e-4)
 
+
+def test_right_padded_cached_batch() -> None:
+    torch.manual_seed(10)
+    model = QyrouArchForCausalLM(tiny_config()).eval()
+    tokens = torch.tensor([[10, 11, 12, 2, 2], [20, 21, 22, 23, 24]])
+    mask = tokens.ne(2).long()
+    cache = QyrouArchHybridCache(model.config.num_hidden_layers)
+    with torch.no_grad():
+        model(tokens, attention_mask=mask, past_key_values=cache, use_cache=True)
+        next_tokens = torch.tensor([[13], [25]])
+        extended_mask = torch.cat((mask, torch.ones((2, 1), dtype=mask.dtype)), dim=1)
+        cached = model(
+            next_tokens,
+            attention_mask=extended_mask,
+            past_key_values=cache,
+            cache_position=torch.tensor([5]),
+            use_cache=True,
+        ).logits[:, -1]
+        full = model(
+            torch.cat((tokens, next_tokens), dim=1),
+            attention_mask=extended_mask,
+            use_cache=False,
+        ).logits[:, -1]
+    torch.testing.assert_close(cached, full, atol=2e-5, rtol=2e-4)
+
+
+def test_unequal_prompt_lengths_match_full_forward() -> None:
+    torch.manual_seed(11)
+    model = QyrouArchForCausalLM(tiny_config()).eval()
+    tokens = torch.tensor([[2, 2, 10, 11, 12, 13], [20, 21, 22, 2, 2, 2]])
+    mask = tokens.ne(2).long()
+    cache = QyrouArchHybridCache(model.config.num_hidden_layers)
+    with torch.no_grad():
+        model(tokens, attention_mask=mask, past_key_values=cache, use_cache=True)
+        next_tokens = torch.tensor([[14], [23]])
+        extended_mask = torch.cat((mask, torch.ones((2, 1), dtype=mask.dtype)), dim=1)
+        cached = model(
+            next_tokens,
+            attention_mask=extended_mask,
+            past_key_values=cache,
+            cache_position=torch.tensor([6]),
+            use_cache=True,
+        ).logits[:, -1]
+        full = model(
+            torch.cat((tokens, next_tokens), dim=1),
+            attention_mask=extended_mask,
+            use_cache=False,
+        ).logits[:, -1]
+    torch.testing.assert_close(cached, full, atol=2e-5, rtol=2e-4)
+
+
+def test_unsupported_arguments_are_rejected() -> None:
+    model = QyrouArchForCausalLM(tiny_config()).eval()
+    tokens = torch.randint(5, 128, (1, 8))
+    with pytest.raises(ValueError, match="Unsupported model arguments"):
+        model(tokens, use_cache=False, made_up_argument=True)
+
+
+def test_generation_with_framework_default_cache_matches_no_cache() -> None:
+    torch.manual_seed(12)
+    model = QyrouArchForCausalLM(tiny_config()).eval()
+    prompt = torch.tensor([[10, 11, 12, 13]])
+    kwargs = {"max_new_tokens": 8, "do_sample": False, "pad_token_id": 2, "eos_token_id": 1}
+    with torch.no_grad():
+        cached = model.generate(prompt, use_cache=True, **kwargs)
+        uncached = model.generate(prompt, use_cache=False, **kwargs)
+    assert cached.shape[1] == 12
+    assert torch.equal(cached, uncached)
+
+
+def test_generation_rejects_unknown_cache_type() -> None:
+    model = QyrouArchForCausalLM(tiny_config()).eval()
+    tokens = torch.randint(5, 128, (1, 8))
+    with pytest.raises(TypeError, match="framework Cache"):
+        model(tokens, past_key_values=object(), use_cache=True)
+

source/tests/test_reference.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("transformers")
+
+from qyrou_arch.configuration_qyrou_arch import QyrouArchConfig
+from qyrou_arch.modeling_qyrou_arch import (
+    CausalConvMixer,
+    CausalGQA,
+    RMSNorm,
+    SwiGLU,
+    QyrouArchForCausalLM,
+    _apply_rope,
+)
+
+torch.manual_seed(3)
+
+
+def reference_config(**overrides) -> QyrouArchConfig:
+    values = dict(
+        vocab_size=128,
+        hidden_size=64,
+        num_hidden_layers=4,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        intermediate_size=128,
+        conv_layers=[2],
+        conv_kernel_size=3,
+        max_position_embeddings=32,
+        qk_norm=False,
+        liger_rms_norm=False,
+        liger_swiglu=False,
+        cut_cross_entropy=False,
+        attention_backend="math",
+        fused_qkv_projection=True,
+        fused_gate_up_projection=True,
+    )
+    values.update(overrides)
+    return QyrouArchConfig(**values)
+
+
+def test_rmsnorm_matches_manual_reference() -> None:
+    module = RMSNorm(64, eps=1e-5)
+    inputs = torch.randn(3, 9, 64)
+    expected = (
+        inputs.float() * torch.rsqrt(inputs.float().square().mean(-1, keepdim=True) + 1e-5)
+    ) * module.weight
+    torch.testing.assert_close(module(inputs), expected, rtol=1e-5, atol=1e-5)
+
+
+def test_rope_matches_manual_reference() -> None:
+    inputs = torch.randn(2, 5, 8)
+    positions = torch.tensor([3, 7, 11, 15, 19], dtype=torch.float32)
+    inv_freq = 1.0 / (10000.0 ** (torch.arange(0, 8, 2).float() / 8))
+    angles = torch.outer(positions, inv_freq)
+    cos, sin = angles.cos(), angles.sin()
+    even, odd = inputs[..., 0::2], inputs[..., 1::2]
+    expected = torch.stack(
+        (even * cos - odd * sin, even * sin + odd * cos), dim=-1
+    ).flatten(-2)
+    torch.testing.assert_close(
+        _apply_rope(inputs, cos.unsqueeze(0), sin.unsqueeze(0)),
+        expected,
+        rtol=1e-5,
+        atol=1e-5,
+    )
+
+
+def test_attention_matches_manual_reference() -> None:
+    config = reference_config()
+    layer = CausalGQA(config, layer_idx=0).eval()
+    inputs = torch.randn(2, 7, 64)
+    positions = torch.arange(7, dtype=torch.float32)
+    inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, 16, 2).float() / 16))
+    angles = torch.outer(positions, inv_freq)
+    cos = angles.cos().unsqueeze(0).unsqueeze(1)
+    sin = angles.sin().unsqueeze(0).unsqueeze(1)
+
+    qkv = layer.qkv_proj(inputs)
+    q, k, v = qkv.split((layer.q_size, layer.kv_size, layer.kv_size), dim=-1)
+    q = q.view(2, 7, 4, 16).transpose(1, 2)
+    k = k.view(2, 7, 2, 16).transpose(1, 2)
+    v = v.view(2, 7, 2, 16).transpose(1, 2)
+    q = _apply_rope(q, cos, sin)
+    k = _apply_rope(k, cos, sin)
+    k = k.repeat_interleave(2, dim=1)
+    v = v.repeat_interleave(2, dim=1)
+    scores = torch.matmul(q, k.transpose(-2, -1)) / 4.0
+    mask = torch.tril(torch.ones((1, 1, 7, 7), dtype=torch.bool))
+    scores = scores.masked_fill(~mask, float("-inf"))
+    weights = torch.softmax(scores, dim=-1)
+    expected = torch.matmul(weights, v).transpose(1, 2).contiguous().view(2, 7, -1)
+    expected = layer.o_proj(expected)
+
+    with torch.no_grad():
+        actual, _ = layer(
+            inputs,
+            (cos, sin),
+            attention_mask=None,
+            cache=None,
+            cache_position=torch.arange(7),
+        )
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
+
+
+def test_convolution_matches_manual_reference() -> None:
+    config = reference_config()
+    layer = CausalConvMixer(config, layer_idx=2).eval()
+    inputs = torch.randn(2, 9, 64)
+    padded = torch.nn.functional.pad(inputs.transpose(1, 2), (2, 0))
+    expected = layer.pointwise(layer.depthwise(padded).transpose(1, 2))
+
+    with torch.no_grad():
+        actual, _ = layer(
+            inputs,
+            (None, None),
+            attention_mask=None,
+            cache=None,
+            cache_position=torch.arange(9),
+        )
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_swiglu_matches_manual_reference() -> None:
+    config = reference_config()
+    module = SwiGLU(config).eval()
+    inputs = torch.randn(2, 9, 64)
+    packed = module.gate_up_proj(inputs)
+    gate, up = packed.chunk(2, dim=-1)
+    expected = module.down_proj(torch.nn.functional.silu(gate) * up)
+
+    with torch.no_grad():
+        actual = module(inputs)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+def test_shifted_loss_matches_manual_cross_entropy() -> None:
+    model = QyrouArchForCausalLM(reference_config()).eval()
+    tokens = torch.randint(5, 128, (2, 8))
+    with torch.no_grad():
+        output = model(tokens, labels=tokens, return_logits=True, use_cache=False)
+        logits = output.logits[:, :-1].float().reshape(-1, 128)
+        targets = tokens[:, 1:].contiguous().view(-1)
+        expected = torch.nn.functional.cross_entropy(logits, targets)
+    torch.testing.assert_close(output.loss, expected, rtol=1e-4, atol=1e-4)
+
+
+def test_attention_weights_sum_to_one_when_requested() -> None:
+    model = QyrouArchForCausalLM(reference_config()).eval()
+    tokens = torch.randint(5, 128, (1, 6))
+    with torch.no_grad():
+        output = model(tokens, use_cache=False, output_attentions=True)
+    assert output.attentions is not None
+    assert len(output.attentions) == 4
+    attention_weights = [weights for weights in output.attentions if weights is not None]
+    assert len(attention_weights) == 3
+    for weights in attention_weights:
+        assert weights.shape[1] == 4
+        torch.testing.assert_close(
+            weights.sum(dim=-1),
+            torch.ones_like(weights.sum(dim=-1)),
+            rtol=1e-4,
+            atol=1e-4,
+        )
+
+
+def test_hidden_states_include_embeddings_and_every_layer() -> None:
+    model = QyrouArchForCausalLM(reference_config()).eval()
+    tokens = torch.randint(5, 128, (1, 6))
+    with torch.no_grad():
+        output = model(tokens, use_cache=False, output_hidden_states=True)
+    assert output.hidden_states is not None
+    assert len(output.hidden_states) == 5
+    assert output.hidden_states[0].shape == (1, 6, 64)
+    assert output.hidden_states[-1].shape == (1, 6, 64)
+    for hidden in output.hidden_states:
+        assert torch.isfinite(hidden).all()

source/tests/test_triton_kernel.py

@@ -4,21 +4,87 @@ import pytest
 
 torch = pytest.importorskip("torch")
 
+cuda_available = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_packed_swiglu_forward_backward_parity() -> None:
-    pytest.importorskip("triton")
+
+def reference_swiglu(packed: torch.Tensor) -> torch.Tensor:
+    gate, up = packed.chunk(2, dim=-1)
+    return torch.nn.functional.silu(gate) * up
+
+
+def assert_kernel_parity(packed: torch.Tensor, atol: float, rtol: float) -> None:
     from qyrou_arch.triton_kernels import PackedSwiGLUFunction
 
+    candidate = packed.detach().clone().requires_grad_(True)
+    output = PackedSwiGLUFunction.apply(candidate)
+    expected = reference_swiglu(packed.detach().requires_grad_(True))
+    gradient = torch.randn_like(output)
+    output.backward(gradient)
+    expected.backward(gradient)
+    torch.testing.assert_close(output, expected, atol=atol, rtol=rtol)
+    torch.testing.assert_close(candidate.grad, packed.grad, atol=atol * 2, rtol=rtol * 2)
+
+
+@cuda_available
+def test_packed_swiglu_forward_backward_parity() -> None:
+    pytest.importorskip("triton")
     torch.manual_seed(11)
     packed = torch.randn(7, 2 * 1328, device="cuda", dtype=torch.bfloat16, requires_grad=True)
     reference = packed.detach().clone().requires_grad_(True)
+    from qyrou_arch.triton_kernels import PackedSwiGLUFunction
+
     output = PackedSwiGLUFunction.apply(packed)
-    gate, up = reference.chunk(2, dim=-1)
-    expected = torch.nn.functional.silu(gate) * up
+    expected = reference_swiglu(reference)
     gradient = torch.randn_like(output)
     output.backward(gradient)
     expected.backward(gradient)
     torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
     torch.testing.assert_close(packed.grad, reference.grad, atol=3e-2, rtol=3e-2)
 
+
+@cuda_available
+def test_packed_swiglu_fp32_parity() -> None:
+    pytest.importorskip("triton")
+    torch.manual_seed(21)
+    packed = torch.randn(11, 2 * 1328, device="cuda", dtype=torch.float32, requires_grad=True)
+    assert_kernel_parity(packed, atol=1e-4, rtol=1e-4)
+
+
+@cuda_available
+@pytest.mark.parametrize("n_cols", [3, 129, 1329])
+def test_packed_swiglu_odd_tail_width(n_cols: int) -> None:
+    pytest.importorskip("triton")
+    torch.manual_seed(31)
+    packed = torch.randn(5, 2 * n_cols, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    assert_kernel_parity(packed, atol=2e-2, rtol=2e-2)
+
+
+@cuda_available
+def test_packed_swiglu_non_contiguous_input() -> None:
+    pytest.importorskip("triton")
+    torch.manual_seed(41)
+    base = torch.randn(7, 2 * 1328 + 4, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    packed = base[:, 2:-2]
+    assert packed.shape[-1] == 2 * 1328
+    assert not packed.is_contiguous()
+    assert_kernel_parity(packed, atol=2e-2, rtol=2e-2)
+
+
+@cuda_available
+def test_packed_swiglu_zero_size_rows() -> None:
+    pytest.importorskip("triton")
+    from qyrou_arch.triton_kernels import PackedSwiGLUFunction
+
+    packed = torch.empty(0, 2 * 1328, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    output = PackedSwiGLUFunction.apply(packed)
+    assert output.shape == (0, 1328)
+    output.sum().backward()
+    assert packed.grad is not None and packed.grad.shape == packed.shape
+
+
+@cuda_available
+def test_packed_swiglu_large_row_count() -> None:
+    pytest.importorskip("triton")
+    torch.manual_seed(51)
+    packed = torch.randn(4096, 2 * 1328, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    assert_kernel_parity(packed, atol=2e-2, rtol=2e-2)