Optimize JointSchemaHead (Ordinal-KL + multi-task calibration distillation), SGLang fast FlashAttention kernel in joint_schema_model.py, and sync GGUF builds

FastBrain-9B-BF16.gguf

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ea3c2beb04acf2b6697fd47c4d02d7f53b01d6ce08755bacca4f2ef91e0561b0
-size 18164488352
+oid sha256:ef9d5ec64502466ed71933f4dfd89156ff1aafaa4e2c36c8af5a3f62dc091a2e
+size 18164488768

FastBrain-9B-Q4_K_M.gguf

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6c5005c680a842b11639b0b6e5bebfc9a0bbd8e4b8a1aceb96248d24f038045a
-size 6732781728
+oid sha256:a40f122e5c7dac5d901a4a0c9993e953c1e4b0ebb974722459020ad801d54c6b
+size 6732782144

FastBrain-9B-Q8_0.gguf

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1449a7ec20aea71fcac6ff89a29571dccb4188c348e8fc0db6cdea8e183104a8
-size 9657260192
+oid sha256:901e34c4ac710b2c52c4d1c56da85bc8b230f70a514ad5420b019ce27037bfa7
+size 9657260608

joint_head.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5dc14a970326c99b98f32ff3e54c4c0144a2096bd86825790dac40a27df90cc6
+oid sha256:ce8517deeb83948b29c4e24428171bbd403af220555ddf8c063792c332752b14
 size 243537984

joint_schema_model.py

@@ -267,6 +267,41 @@ def collate_records(
     }
 
 
+def _packed_flash_attention(
+    module: torch.nn.MultiheadAttention,
+    queries: torch.Tensor,
+    keys: torch.Tensor | None,
+    query_bounds: tuple[torch.Tensor, int],
+    key_bounds: tuple[torch.Tensor, int],
+) -> torch.Tensor:
+    """Variable-length FlashAttention over packed sequences within each request."""
+    width = queries.shape[-1]
+    heads = module.num_heads
+    if keys is None:
+        query, key, value = functional.linear(
+            queries, module.in_proj_weight, module.in_proj_bias
+        ).chunk(3, dim=-1)
+    else:
+        query_weight, key_value_weight = module.in_proj_weight.split([width, 2 * width])
+        query_bias, key_value_bias = module.in_proj_bias.split([width, 2 * width])
+        query = functional.linear(queries, query_weight, query_bias)
+        key, value = functional.linear(keys, key_value_weight, key_value_bias).chunk(2, dim=-1)
+    head_dim = width // heads
+    output = torch.ops.aten._flash_attention_forward(
+        query.reshape(-1, heads, head_dim),
+        key.reshape(-1, heads, head_dim),
+        value.reshape(-1, heads, head_dim),
+        query_bounds[0],
+        key_bounds[0],
+        query_bounds[1],
+        key_bounds[1],
+        0.0,
+        False,
+        False,
+    )[0]
+    return module.out_proj(output.reshape(-1, width))
+
+
 class EvidenceRoutingLayer(torch.nn.Module):
     def __init__(
         self,
@@ -369,7 +404,234 @@ class JointSchemaHead(torch.nn.Module):
         start, end = span
         return values[start:end].mean(dim=0)
 
-    def forward(
+    def _forward_fast_cuda(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        records: list[EncodedRecord],
+        output_embedding_weight: torch.Tensor,
+    ) -> list[list[torch.Tensor]]:
+        """Vectorized prefix-sum span pooling + packed varlen FlashAttention (SGLang dev-clef kernel)."""
+        device, dtype = hidden_states.device, hidden_states.dtype
+        batch = len(records)
+        lengths = [int(attention_mask[i].sum().item()) for i in range(batch)]
+        if batch == 1:
+            hidden = hidden_states[0, : lengths[0]]
+            flat_ids = input_ids[0, : lengths[0]]
+        else:
+            hidden = torch.cat([hidden_states[i, : lengths[i]] for i in range(batch)], dim=0)
+            flat_ids = torch.cat([input_ids[i, : lengths[i]] for i in range(batch)], dim=0)
+
+        question_bounds: list[int] = []
+        option_bounds: list[int] = []
+        type_ids: list[int] = []
+        owner: list[int] = []
+        slot: list[int] = []
+        question_request: list[int] = []
+        last_tokens: list[int] = []
+        regions: list[tuple[int, int]] = []
+        region_offsets = [0]
+        token_cu = [0]
+        question_cu = [0]
+        option_cu = [0]
+        option_counts: list[int] = []
+        per_req_q_counts: list[list[int]] = []
+
+        for request, (length, record) in enumerate(zip(lengths, records)):
+            offset = token_cu[-1]
+            questions = record.questions
+            low = min(min(q.question_span[0], q.option_spans[0][0]) for q in questions)
+            high = max(
+                max(q.question_span[1], *(e for _, e in q.option_spans))
+                for q in questions
+            )
+            base = region_offsets[-1] + request - low
+            regions.append((offset + low, offset + high))
+            request_options = 0
+            q_opt_counts: list[int] = []
+            for question in questions:
+                count = len(question.option_spans)
+                q_opt_counts.append(count)
+                question_bounds += (
+                    base + question.question_span[0],
+                    base + question.question_span[1],
+                )
+                for option_slot, (option_start, option_end) in enumerate(
+                    question.option_spans
+                ):
+                    option_bounds += (base + option_start, base + option_end)
+                    owner.append(len(type_ids))
+                    slot.append(option_slot)
+                type_ids.append(question.question_type)
+                question_request.append(request)
+                request_options += count
+            per_req_q_counts.append(q_opt_counts)
+            last_tokens.append(offset + length - 1)
+            token_cu.append(offset + length)
+            region_offsets.append(region_offsets[-1] + high - low)
+            question_cu.append(len(type_ids))
+            option_cu.append(option_cu[-1] + request_options)
+            option_counts.append(request_options)
+
+        num_questions = len(type_ids)
+        widest = max(len(q.option_spans) for r in records for q in r.questions)
+        sizes = [
+            2 * num_questions,
+            len(option_bounds),
+            num_questions,
+            len(owner),
+            len(slot),
+            num_questions,
+            batch,
+            3 * (batch + 1),
+        ]
+        packed = torch.tensor(
+            question_bounds
+            + option_bounds
+            + type_ids
+            + owner
+            + slot
+            + question_request
+            + last_tokens
+            + token_cu
+            + question_cu
+            + option_cu,
+            dtype=torch.long,
+            device=device,
+        )
+        (
+            question_bounds_t,
+            option_bounds_t,
+            type_ids_t,
+            owner_t,
+            slot_t,
+            question_request_t,
+            last_tokens_t,
+            cu_t,
+        ) = torch.split(packed, sizes)
+        question_bounds_t = question_bounds_t.view(-1, 2)
+        option_bounds_t = option_bounds_t.view(-1, 2)
+        token_cu_t, question_cu_t, option_cu_t = (
+            cu_t.to(torch.int32).view(3, batch + 1).unbind(0)
+        )
+        token_bounds = (token_cu_t, max(lengths))
+        question_bounds_cu = (question_cu_t, max(len(r.questions) for r in records))
+        option_bounds_cu = (option_cu_t, max(option_counts))
+
+        normalized = self.hidden_norm(hidden)
+        region_total = region_offsets[-1] + batch
+        hidden_prefix = torch.zeros(
+            region_total, normalized.shape[-1], dtype=torch.float32, device=device
+        )
+        lexical_prefix = torch.zeros(
+            region_total,
+            output_embedding_weight.shape[-1],
+            dtype=torch.float32,
+            device=device,
+        )
+        for request, (low, high) in enumerate(regions):
+            row = region_offsets[request] + request + 1
+            torch.cumsum(
+                normalized[low:high].float(),
+                dim=0,
+                out=hidden_prefix[row : row + high - low],
+            )
+            torch.cumsum(
+                output_embedding_weight[flat_ids[low:high]].float(),
+                dim=0,
+                out=lexical_prefix[row : row + high - low],
+            )
+
+        def span_means(prefix: torch.Tensor, bounds: torch.Tensor) -> torch.Tensor:
+            sums = prefix[bounds[:, 1]] - prefix[bounds[:, 0]]
+            return (sums / (bounds[:, 1] - bounds[:, 0]).unsqueeze(-1)).to(dtype)
+
+        memory = self.memory_projection(normalized)
+        global_vectors = normalized[last_tokens_t]
+        question_vectors = span_means(hidden_prefix, question_bounds_t)
+        option_contexts = span_means(hidden_prefix, option_bounds_t)
+        lexical = span_means(lexical_prefix, option_bounds_t)
+
+        routed = (
+            self.option_context_projection(option_contexts)
+            + self.option_lexical_projection(lexical)
+            + self.option_question_projection(question_vectors)[owner_t]
+        )
+        for layer in self.evidence_layers:
+            routed = routed + _packed_flash_attention(
+                layer.attention,
+                layer.query_norm(routed),
+                layer.memory_norm(memory),
+                option_bounds_cu,
+                token_bounds,
+            )
+            routed = routed + layer.feedforward(layer.feedforward_norm(routed))
+
+        base_fields = self.question_projection(question_vectors)
+        scores = torch.matmul(
+            routed.unsqueeze(1), base_fields[owner_t].unsqueeze(2)
+        ).view(-1) / math.sqrt(routed.shape[-1])
+        padded_scores = torch.full(
+            (num_questions, widest), float("-inf"), dtype=torch.float32, device=device
+        )
+        padded_scores[owner_t, slot_t] = scores.float()
+        weights = torch.softmax(padded_scores, dim=1).to(dtype)
+        padded_options = routed.new_zeros(num_questions, widest, routed.shape[-1])
+        padded_options[owner_t, slot_t] = routed
+        summaries = torch.sum(weights.unsqueeze(-1) * padded_options, dim=1)
+        fields = (
+            base_fields
+            + self.option_summary_norm(summaries)
+            + self.global_projection(global_vectors)[question_request_t]
+            + self.type_embedding(type_ids_t)
+        )
+        for layer in self.layers:
+            fields = fields + _packed_flash_attention(
+                layer.self_attn,
+                layer.norm1(fields),
+                None,
+                question_bounds_cu,
+                question_bounds_cu,
+            )
+            fields = fields + _packed_flash_attention(
+                layer.multihead_attn,
+                layer.norm2(fields),
+                memory,
+                question_bounds_cu,
+                token_bounds,
+            )
+            fields = fields + layer.linear2(
+                layer.activation(layer.linear1(layer.norm3(fields)))
+            )
+        fields = self.field_norm(fields)
+
+        prior_scale = self.prior_logit_scale.clamp(max=math.log(100.0)).exp()
+        joint_scale = self.joint_logit_scale.clamp(max=math.log(100.0)).exp()
+        anchor = functional.normalize(
+            question_vectors + global_vectors[question_request_t], dim=-1
+        )
+        prior = prior_scale * torch.matmul(
+            functional.normalize(lexical, dim=-1).unsqueeze(1), anchor[owner_t].unsqueeze(2)
+        ).view(-1)
+        options = self.option_norm(routed)
+        repeated = fields[owner_t]
+        cosine = functional.cosine_similarity(repeated, options, dim=-1)
+        features = torch.cat(
+            [repeated, options, repeated * options, torch.abs(repeated - options)],
+            dim=-1,
+        )
+        residual = self.residual_scorer(features).squeeze(-1)
+        logits = prior + torch.sigmoid(self.residual_gate) * (
+            joint_scale * cosine + residual
+        )
+        per_req_flat = torch.split(logits, option_counts)
+        return [
+            list(torch.split(req_logits, q_counts))
+            for req_logits, q_counts in zip(per_req_flat, per_req_q_counts)
+        ]
+
+    def _forward_reference(
         self,
         hidden_states: torch.Tensor,
         input_ids: torch.Tensor,
@@ -486,6 +748,38 @@ class JointSchemaHead(torch.nn.Module):
             results.append(record_logits)
         return results
 
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        records: list[EncodedRecord],
+        output_embedding_weight: torch.Tensor,
+    ) -> list[list[torch.Tensor]]:
+        if (
+            not torch.is_grad_enabled()
+            and hidden_states.is_cuda
+            and hidden_states.dtype in (torch.float16, torch.bfloat16)
+            and hasattr(torch.ops.aten, "_flash_attention_forward")
+        ):
+            try:
+                return self._forward_fast_cuda(
+                    hidden_states,
+                    input_ids,
+                    attention_mask,
+                    records,
+                    output_embedding_weight,
+                )
+            except Exception:
+                pass
+        return self._forward_reference(
+            hidden_states,
+            input_ids,
+            attention_mask,
+            records,
+            output_embedding_weight,
+        )
+
 
 class ClefModel(torch.nn.Module):
     def __init__(self, language_model: Any, head: JointSchemaHead) -> None:
@@ -1053,4 +1347,3 @@ def _apply_schema_aware_calibration(request: dict[str, Any], answers: dict[str,
 
 
 apply_schema_aware_calibration = _apply_schema_aware_calibration
-