Support Cohere Command-A (Cohere2ForCausalLM arch)

[aikitoria]

It would be great to support this new model! https://cohere.com/blog/command-a

They use a fairly unique architecture, where some layers use sliding window attention while others use global attention with no position embeddings, so even though I read through the documentation on how to add a model I'm a little lost on how to do this myself.

[user-0a]

+1, would also like support for this!

[aikitoria]

I've just finished implementing this, it appears to be working! I've also implemented support for FP8 Rowwise quants for the Cohere models. This was only missing a small piece where FP8 support was implemented for RmsNorm CUDA kernels but not LayerNorm ones.

Will open a PR once I ran a few more tests.

[aikitoria]

Hmm, there are problems with kv cache reuse. If I generate something that exceeds the sliding window size, and then generate something else that shares a prefix, it generates total garbage. Perhaps the implementation of different sliding windows per layer is not compatible with paged cache and cache reuse.

Luckily, nvidia recently open-sourced executor and batch_manager, which means we can just... fix it!

[aikitoria]

I've put the WIP branch here, it is not completely finished (i.e. does not yet load sliding window layers from config) but it can compile and run the model successfully: https://github.com/aikitoria/TensorRT-LLM/tree/experiments

Still haven't figured out how to fix the KV cache reuse with the sliding window though. And the changes to the LayernormQuantization plugin might have broken other models

Right now you have to add enable_block_reuse=False, max_attention_window=[4096, 4096, 4096, 131072] to KvCacheConfig for the model to not produce garbage, which is terrible, I will try to fix it but this part of the code is very complex and the iteration time is absolutely horrible due to huge templates

[aikitoria]

Ok, after looking into it more, I'm pretty sure it's IMPOSSIBLE for us to fix this. To allow block reuse to work, we need to not use a cyclic kv cache, instead we need to keep kv cache for the entire sequence, and only limit the window considered during attention. However, trtllm currently always uses cyclic kv cache for each layer where a sliding window is enabled, even though this is not actually useful at all for this model as it also has layers without sliding windows. The early kv blocks being overwritten for those layers then corrupts the blocks for reusing them across different requests*.

Fixing this would be possible if we could edit the kernels. However we are missing the code for the context phase kernel, there is only cubin stuff here: https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention

So there's nothing we can do.

@byshiue can you guys please either release the code so we can fix it or officially support Cohere2ForCausalLM including block reuse (and the full sequence length, not just one sliding window like was done for Gemma 2)? The model runs fine in vLLM but trtllm runs this model more than two times faster so this is really frustrating that we can't make it work fully.

*I might still be misunderstanding the exact issue, but it's difficult to investigate further when there is no code.

[ming-wei]

Thanks for raising the issue.

We are aware of the garbage result issue when kv cache reuse and sliding window attention both enabled. We are on it right now.

[aikitoria]

@ming-wei That would be great, thank you!

From exploring batch_manager code, it looks like someone had the idea of solving it by disabling block reuse if a sliding window is used, but it isn't working because it compares to the max num tokens rather than the smallest window size. Regardless, please don't let that become the solution!

It could fully support block reuse on this model if only there was a way to enable sliding windows without using cyclic kv cache at the same time. Adding this to the generation stage kernel was pretty easy in the end (keep the loop as it was, but add window start offset to timeslots when accessing kv) but I can't modify context one because it's missing the code.

[zhaocc1106]

It's exciting for your work for commanda. But i found some similar promblem for sliding window attention when support gemma3 text llm. I found attention mask of sliding window attention is not a pure causal type, which result in error result for long input (longer than sliding window) if i use AttentionMaskType.causal type . Especially for gemma3 with sliding window only 512.

[aikitoria]

It is difficult to tell exactly what goes wrong in the context kernel as we don't have the code, it seemed to me like it is actually working properly when passed attention_mask_type = ContextAttentionMaskType::SLIDING_WINDOW_CAUSAL to fmha params in combination with cyclic kv cache. But I don't want cyclic kv cache so block reuse works.

[zhaocc1106]

Thanks for answering. But have you tested with long input which >4096(the window)? It’s easy to found the issue because gemma3 window only 512. I test gemm2 with long input which > 4096 (the window). It seems the phenomenon is the same.

[aikitoria]

I tried some longer inputs and they were working but I didn't do super intensive testing yet

[zhaocc1106]

How do you set the attention mask to sliding window? I found the attention_mask_type you set is casual not sliding window causal, as following:
https://github.com/aikitoria/TensorRT-LLM/blob/cc798be8f9f672ab7751cbc38deb94b5bdad8e64/tensorrt_llm/models/commanda/model.py#L43
I found the AttentionMaskType.sliding_window_causal not be supported properly as comment: #2880 (comment)

Thanks~

[aikitoria]

Good point, that should probably be fixed in my code, but It automatically changes the mask type to sliding window here so it should be fine https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp#L300

And the generation phase kernel does not have this mask type argument (it is purely controlled by the sliding window size)