[Cookbook] Nemotron3-Ultra: Add mamba-backend and SSM dtype flags (#28675)
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai>
This commit is contained in:
co-authored by
Brayden Zhong
parent
ab21dc984a
commit
34e5e38604
@@ -60,12 +60,20 @@ The generator only emits a runnable command for combinations that NVIDIA / SGLan
|
||||
- **Attention backend**:
|
||||
|
||||
**H100/H200**: Use flash attention 3 backend by default.
|
||||
**B200/GB200/B300/GB300**: Append `--attention-backend trtllm_mha`. The flashinfer default breaks the overlap scheduler on Blackwell, so `trtllm_mha` is required there.
|
||||
**B200/GB200/B300/GB300**: Set `--attention-backend trtllm_mha`. The flashinfer default breaks the overlap scheduler on Blackwell, so `trtllm_mha` is required there.
|
||||
|
||||
- **Mamba scheduler strategy**:
|
||||
|
||||
Always launch with `--mamba-scheduler-strategy extra_buffer`. This hybrid Transformer-Mamba model requires the `extra_buffer` strategy for correct scheduling of its Mamba state.
|
||||
|
||||
- **Mamba backend**:
|
||||
|
||||
The Mamba layers use the Triton SSM kernels by default. For better performance, set `--mamba-backend flashinfer` to use the FlashInfer Mamba kernels instead.
|
||||
|
||||
- **Mamba SSM precision**:
|
||||
|
||||
The SSM state dtype defaults to the model config value. Set `--mamba-ssm-dtype float16` to store the Mamba states in FP16, which reduces mamba cache memory without significant accuracy loss.
|
||||
|
||||
- **TP support**:
|
||||
|
||||
To set tp size, use `--tp <4|8|16>`. Recommended pairings:
|
||||
@@ -82,7 +90,7 @@ The generator only emits a runnable command for combinations that NVIDIA / SGLan
|
||||
|
||||
- **Expert parallel (EP)**:
|
||||
|
||||
This MoE only supports `ep_size == 1` (off) or `ep_size == tp_size`. To enable expert parallelism, append `--ep <tp>` with the same value as `--tp`.
|
||||
This MoE only supports `ep_size == 1` (off) or `ep_size == tp_size`. To enable expert parallelism, set `--ep <tp>` with the same value as `--tp`.
|
||||
|
||||
- **Multi-token prediction (MTP)**:
|
||||
|
||||
@@ -90,15 +98,15 @@ The generator only emits a runnable command for combinations that NVIDIA / SGLan
|
||||
|
||||
- **FP8 KV cache**:
|
||||
|
||||
To enable fp8 kv cache, please append `--kv-cache-dtype fp8_e4m3`.
|
||||
To enable fp8 kv cache, set `--kv-cache-dtype fp8_e4m3`. This is enabled by default on the NVFP4 checkpoint.
|
||||
|
||||
- **Reasoning parser**:
|
||||
|
||||
Append `--reasoning-parser nemotron_3` to enable structured reasoning traces (`reasoning_content` field in the response).
|
||||
Set `--reasoning-parser nemotron_3` to enable structured reasoning traces (`reasoning_content` field in the response).
|
||||
|
||||
- **Tool calling**:
|
||||
|
||||
Append `--tool-call-parser qwen3_coder` to enable tool calling support.
|
||||
Set `--tool-call-parser qwen3_coder` to enable tool calling support.
|
||||
|
||||
## 4. Model Invocation
|
||||
|
||||
|
||||
@@ -154,6 +154,24 @@ export const Nemotron3UltraDeployment = () => {
|
||||
{ id: 'bf16', label: 'bf16', default: false }
|
||||
]
|
||||
},
|
||||
mambabackend: {
|
||||
name: 'mambabackend',
|
||||
title: 'Mamba Backend',
|
||||
items: [
|
||||
{ id: 'triton', label: 'Triton', subtitle: 'Default', default: true },
|
||||
{ id: 'flashinfer', label: 'FlashInfer', subtitle: 'Faster', default: false }
|
||||
],
|
||||
commandRule: (value) => value === 'flashinfer' ? '--mamba-backend flashinfer' : null
|
||||
},
|
||||
mambassmdtype: {
|
||||
name: 'mambassmdtype',
|
||||
title: 'Mamba SSM DType',
|
||||
items: [
|
||||
{ id: 'default', label: 'Default', subtitle: 'Model config', default: true },
|
||||
{ id: 'float16', label: 'float16', subtitle: 'Less memory', default: false }
|
||||
],
|
||||
commandRule: (value) => value === 'float16' ? '--mamba-ssm-dtype float16' : null
|
||||
},
|
||||
thinking: {
|
||||
name: 'thinking',
|
||||
title: 'Reasoning Parser',
|
||||
|
||||
Reference in New Issue
Block a user