This bounty is for bringing up CohereForAI/c4ai-command-r-v01 using TTNN APIs on Tenstorrent hardware (Wormhole T3K — 4× N300s).
Command R is Cohere’s open-weight 35B parameter autoregressive language model built on an optimized transformer architecture with 128K context length. It was specifically designed for grounded generation and retrieval-augmented generation (RAG), with native single-step and multi-step tool use. The architecture closely follows Llama-style decoder-only transformers with GQA, RoPE positional embeddings, SwiGLU activations, and a unique QK LayerNorm applied within each attention block.
At 35B parameters, this model is a natural fit for a T3K (4× N300, TP=8) topology and represents a practical bringup target: its operator composition closely matches what is already proven on TT hardware (Llama, Qwen, Mixtral), while adding new coverage for the Cohere model type, QK LayerNorm, and grounded-generation prompt templates.
What Success Looks Like
Implement the full Command R decoder stack in TTNN: token embeddings, 40× CohereDecoderLayer (QK-normed GQA attention + SwiGLU MLP), and LM head
Run decode-mode inference (autoregressive generation) on a T3K mesh with TP=8 across chips using TT-Fabric CCL ops
Produce text outputs that match the HuggingFace reference numerically (PCC ≥ 0.99 per layer)
Validate both chat and grounded-generation (RAG citation) prompt templates
Measure and report throughput (T/S/U) around 60/s on T3K and time-to-first-token (TTFT) around 70 ms on device.
Stage 1 — Bring-Up
Implement c4ai-command-r-v01 using TTNN APIs (Python) targeting a WH T3K mesh (4× N300, 8 chips, TP=8)
Model runs end-to-end on T3K with no errors
All 40 CohereDecoderLayer blocks run in sequence; validate each submodule’s output (attention, MLP, QK LayerNorm) against the HF PyTorch reference using PCC
Produces valid text tokens on sample prompts; verify logits against the HF reference at decode step 1
Stage 2 — Basic Optimizations
Use optimal sharded/interleaved memory configs across all 40 decoder layers
Efficient tensor-parallel sharding for Q/K/V/O projections (TP=8) and MLP gate/up/down projections
Distribute KV cache across 8 devices using TT-Fabric CCL (AllReduce for attention output, AllGather for MLP)
Fuse QK LayerNorm + scaling into the attention kernel where possible
Store KV cache and intermediate activations in L1 where beneficial
Achieve initial decode throughput within order-of-magnitude of target
Stage 3 — Deeper Optimization
Maximize core utilization per decode step across all 8 chips
Apply BFP8/BFP4 weight compression for MLP FF1/FF3 layers to improve memory bandwidth utilization
Minimize tensor reshape and layout-conversion overheads between layers
Enable chunked prefill for long-context RAG use cases (up to 128K tokens)
Meet or exceed target decode throughput (T/S/U) at batch size 32
Guidance & Starting Points
Parameter
Value
Parameters
35B
Layers
40
Hidden size
8192
Intermediate size
22528
Attention heads (Q)
64
KV heads (GQA)
8
Context length
128K
Positional encoding
RoPE
Activation
SwiGLU
Special ops
QK LayerNorm per attention block
Key differences from Llama (things to watch during bringup)
CohereModel uses a query_key_layer_scaling_coeff and applies LayerNorm directly on Q and K tensors inside each attention block — this is not present in Llama/Mistral and will require a small TTNN kernel sequence (layernorm → scale → matmul)
The LM head uses a logit_scale factor applied post-linear.
Possible Approaches
Start from the HuggingFace CohereForCausalLM / PyTorch reference and port layers one by one to TTNN, validating each submodule’s output against the reference before full integration. The models/tt_transformers directory in tt-metal is the recommended starting point — CohereDecoderLayer maps closely to the existing TtTransformerBlock pattern used by Llama and Qwen.
Re-use the existing GQA attention kernel from Llama/Qwen and add a post-projection QK LayerNorm step; validate PCC against HF before proceeding to the MLP.
Use TTNN profiling tools (tt-exalens, perf report headers) to identify bottlenecks and fusion opportunities after Stage 1 correctness is established.
For the T3K mesh setup, follow the T3K mesh device configuration used in existing QuietBox model demos (models/demos/t3000/).
Result Submission Guidelines
Beyond the model implementation itself, contributors must submit the following material as proof of work.
However, feel free to open a PR at any time if you want us checking that you are on the right track.
Just understand that payout is only made after all 3 stages are completed.
Background
What Success Looks Like
Guidance & Starting Points
Possible Approaches
Result Submission Guidelines
Resources