Implementation:Predibase Lorax Flash Cohere Modeling

Knowledge Sources	Predibase_Lorax
Domains	Model_Architecture, Inference
Last Updated	2026-02-08 00:00 GMT

Overview

Optimized Cohere Command transformer implementation for LoRax inference serving with flash attention, custom rotary embeddings, per-head layer normalization, and LoRA adapter support.

Description

FlashCohereForCausalLM implements the Cohere Command model architecture with flash attention for efficient batched inference. The module features a custom interleaved rotary embedding scheme and per-head layer normalization that distinguish it from standard Llama-style implementations.

The file contains seven classes organized as a layered architecture:

• CohereRotary -- Custom rotary embedding class extending PositionRotaryEmbedding that applies rotary embeddings using interleaved even/odd element splitting rather than the standard first-half/second-half approach.
• CohereLayerNorm -- Per-head layer normalization that applies separate weight matrices per attention head, using fused dropout_layer_norm kernels for dimensions up to 8192 and a manual implementation for larger dimensions.
• FlashCohereAttention -- Multi-head attention with GQA support, the custom Cohere rotary embeddings, flash attention for prefill and paged attention for decode. Supports adapter-aware Q/K/V/O projections via TensorParallelMultiAdapterLinear.
• CohereMLP -- Gated MLP with fused gate-up projections and adapter-aware gate/up/down projection layers.
• FlashCohereLayer -- Single transformer decoder layer combining attention and MLP with pre-norm layer normalization. Uses the Cohere-specific parallel attention and MLP computation pattern.
• FlashCohereModel -- Full transformer model stacking N decoder layers with token embeddings and final normalization.
• FlashCohereForCausalLM -- Top-level causal language model that wraps the model with a language model head and applies logit_scale to the output logits and speculative logits.

The implementation supports tensor parallelism for multi-GPU serving and uses the rotary_emb native extension for efficient rotary embedding computation.

Usage

Used internally by the LoRax server when serving Cohere Command models. Loaded via the model registry when the model config type matches.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/flash_cohere_modeling.py
• Lines: 1-551

Signature

class FlashCohereForCausalLM(torch.nn.Module):
    def __init__(self, prefix: str, config, weights):
        ...

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlen_prefill: Optional[torch.Tensor],
        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
        block_tables: torch.Tensor,
        slots: torch.Tensor,
        seqlen: Seqlen,
        max_s: int,
        adapter_data: AdapterBatchData,
        prefill_cache_indices: Optional[torch.Tensor] = None,
        lm_head_indices: Optional[torch.Tensor] = None,
        skip_lm_head: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        ...

Import

from lorax_server.models.custom_modeling.flash_cohere_modeling import FlashCohereForCausalLM

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	torch.Tensor	Yes	Token IDs [batch_size, seq_len]
position_ids	torch.Tensor	Yes	Position indices for rotary embeddings
cu_seqlen_prefill	Optional[torch.Tensor]	Yes	Cumulative sequence lengths for flash attention prefill (None during decode)
kv_cache	List[Tuple[torch.Tensor, torch.Tensor]]	Yes	Key-value cache tensors per layer
block_tables	torch.Tensor	Yes	Block table indices for paged attention
slots	torch.Tensor	Yes	Slot indices for KV cache placement
seqlen	Seqlen	Yes	Sequence length metadata wrapper
max_s	int	Yes	Maximum sequence length in the batch
adapter_data	AdapterBatchData	Yes	LoRA adapter weights and indices for the batch
prefill_cache_indices	Optional[torch.Tensor]	No	Indices for selective KV cache population during prefill
lm_head_indices	Optional[torch.Tensor]	No	Indices to select specific positions for LM head output
skip_lm_head	bool	No	If True, return hidden states without applying the LM head

Outputs

Name	Type	Description
logits	torch.Tensor	Next-token logits [batch_size, vocab_size] scaled by logit_scale (or hidden states if skip_lm_head is True)
speculative_logits	Optional[torch.Tensor]	Speculative decoding logits scaled by logit_scale, or None

Usage Examples

# Internal usage within LoRax server
from lorax_server.models.custom_modeling.flash_cohere_modeling import FlashCohereForCausalLM

# Model is instantiated by the model registry, not directly by users
# See server/lorax_server/models/__init__.py for registration

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime