diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 2e234659c34c..4d543f9deee2 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -10,7 +10,15 @@ This guide uses Slurm and the `trtllm-llmapi-launch` multi-node launcher. The co ## Prerequisites -* GPU: NVIDIA Blackwell GPUs. DEP16 and TEP16 use 16 GPUs; the TEP8 recipe uses 8 GPUs. These deployment recipes were validated on GB300 NVL GPUs. The repository's Slurm examples assume 4 GPUs per node. Other GPU architectures are not currently supported. +* GPU: NVIDIA Blackwell GPUs. DEP16 and TEP16 use 16 GPUs; the TEP8 recipe uses 8 GPUs. These deployment recipes were validated on GB300 NVL GPUs. The repository's Slurm examples assume 4 GPUs per node. Per-GPU memory requirements differ per recipe, and are set by the attention layout rather than by the GPU count: + + | Recipe | Attention layout | Per-rank weights | Requires | + | :-- | :-- | --: | :-- | + | DEP16 | attention-DP (replicated) | 210 GB | GB300-class per-GPU memory | + | TEP8 | attention-TP, EP8 | 213 GB | GB300-class per-GPU memory | + | TEP16 | attention-TP, EP16 | 115 GB | validated on GB200 (`SM100`) | + + DEP16 replicates the BF16 non-expert weights on every rank (114 GB per rank) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB per rank). TEP16 shards those non-expert weights instead, which is what brings it within `SM100` per-GPU memory; TEP8 does not fit because its 8-way expert share alone is 181 GB per rank. On B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level and covered by unit tests in CI. Other GPU architectures are not supported. * Multi-node launcher: Slurm with the pyxis/enroot container plugin (or an equivalent MPI launcher) to start one rank per GPU across the nodes. * High-speed inter-node interconnect (e.g., NVLink/InfiniBand) for the expert-parallel traffic. * Shared filesystem visible to all nodes for the repository, the model weights, and the configuration file. @@ -26,7 +34,7 @@ The checkpoint and the configuration file must live on a shared filesystem visib ## Feature Support Notes -* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The configurations and results in this guide were validated on NVIDIA GB300 NVL GPUs. Support for other GPU architectures may be added in a future release. +* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The performance results in this guide were validated on NVIDIA GB300 NVL GPUs. Kimi K3 kernels and modules are also functional on B200 (`SM100`) and covered by unit tests in CI, and the TEP16 deployment is validated end-to-end on GB200 (`SM100`); DEP16 and TEP8 require GB300-class per-GPU memory (see Prerequisites). Support for other GPU architectures may be added in a future release. * **High-throughput and low-latency deployments are provided.** DEP16 (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`) is the high-throughput deployment. TEP16 (`enable_attention_dp: false`, `moe_expert_parallel_size: 16`) is the low-latency deployment. An 8-GPU deployment, TEP8 (`enable_attention_dp: false`, `moe_expert_parallel_size: 8`), is also provided. Select the deployment and concurrency appropriate for your workload. * **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`. * **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps. @@ -44,7 +52,7 @@ python3 scripts/build_wheel.py --cuda_architectures 103-real --skip_building_whe .venv-3.12/bin/python -m pip install -e . ``` -`build_wheel.py` creates the virtual environment at the repository root, named after the container's Python version: `.venv-3.12` for the current containers (Python 3.12). If your container ships a different Python, substitute the matching `.venv-.` path in the commands on this page. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300). The multi-node jobs below run TensorRT LLM from this in-place environment, so build and install with the repository at the same path the jobs use. +`build_wheel.py` creates the virtual environment at the repository root, named after the container's Python version: `.venv-3.12` for the current containers (Python 3.12). If your container ships a different Python, substitute the matching `.venv-.` path in the commands on this page. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300, `100-real` for B200). A `103-real` build also runs on B200 (the Kimi K3 kernels compile for the `100f` family) but omits the `sm100a`-specific batched-GEMM kernels, so build with `100-real` when targeting B200. The multi-node jobs below run TensorRT LLM from this in-place environment, so build and install with the repository at the same path the jobs use. Kimi K3 additionally depends on `fla` and `einops`, installed into the same in-place environment (these dependencies might be removed in future releases, replaced by other kernels): @@ -271,6 +279,34 @@ The job writes progress and results to `kimi-k3-eval-.log` in the submis The expected average accuracy is approximately 96.47. Small differences (roughly ±0.5 points) are possible with different checkpoint or dependency revisions. +### TEP16 on GB200 + +The same job runs the TEP16 layout with `--parallel tep`, which rewrites a per-job copy of the evaluation YAML with `enable_attention_dp: false` and raises `max_batch_size` from 32 to 128 (with attention-DP off every rank serves the same global batch instead of its own, so the batch size is raised to recover eval concurrency): + +```bash +sbatch --account --partition batch --qos --time 04:00:00 \ + examples/kimi_k3/run_eval_kimi_k3.sbatch \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh \ + --task gsm8k --parallel tep +``` + +The batch script declares `--nodes=4 --ntasks-per-node=4 --gpus-per-node=4`, and takes `--account`, `--partition` and `--qos` from the submitting command line. Export `KIMI_K3_ROUTER_BF16=0` before submitting: with attention-DP off the MoE router gate defaults to its BF16 fast path, which can flip borderline expert picks, so the reference scores above are only comparable with that path disabled. `KIMI_K3_FP8_WEIGHT_READ` defaults to `0`, which is the precision the reference scores were measured at. + +Measured on 16 GB200 GPUs (4 nodes, `100-real` build, 184.31 GiB per GPU), with the checkpoint's native MXFP4 routed experts: + +| Filter | Exact match | +| :-- | --: | +| Flexible extract | 96.82 | +| Strict match | 96.74 | + +| Per-rank memory | Value | +| :-- | --: | +| Weights | 106.67 GiB | +| Non-torch (NCCL, CUDA graphs) | 15.32 GiB | +| Peak during profiling | 125.96 GiB | +| KV cache at `free_gpu_memory_fraction: 0.25` | 15.60 GiB | + ## Benchmarking Performance ### Run the End-to-End Performance Sweep diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 7e8fb78e6e3a..1d21af9b0a69 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -32,6 +32,8 @@ The following is a table of supported models for the PyTorch backend: | `HunYuanMoEForCausalLM` [^5] | Hunyuan MoE | `tencent/Hunyuan-A13B-Instruct` | | `InternLM3ForCausalLM` [^5] | InternLM3 | `internlm/internlm3-8b-instruct` | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | `moonshotai/Kimi-K2.5` | +| `KimiK3ForConditionalGeneration` [^15]| Kimi-K3 | `moonshotai/Kimi-K3` | +| `KimiLinearForCausalLM` [^15] | Kimi-K3 (text decoder) | `moonshotai/Kimi-K3` | | `LagunaForCausalLM` | Laguna-XS | `poolside/laguna-XS.2` | | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | | `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | @@ -99,6 +101,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models). [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`). +[^15]: Kimi K3 is only supported on NVIDIA Blackwell GPUs (`SM100` family); see the [Kimi K3 deployment guide](../deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md). Which recipes fit is set by the attention layout rather than by the GPU count: DEP16 (`enable_attention_dp: true`) replicates the BF16 non-expert weights on every rank (114 GB) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB), needing 210 GB per rank, so it requires GB300-class per-GPU memory. TEP16 (`enable_attention_dp: false`) shards those non-expert weights instead and needs 115 GB per rank; it is validated end-to-end on GB200 (`SM100`) at 16 GPUs. On B200 (`SM100`), kernel and module support is functional and covered by CI. Note that the FP8 weight-read path (TRTLLM-14765) does not relax the DEP16 requirement: the conversion runs after the weights are already resident in BF16, so it lowers the steady-state footprint but not the load-time peak. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 14d6023cd713..27f75b8d2b70 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -5,8 +5,15 @@ start and configuration for GSM8K evaluation. ## Hardware support -Only NVIDIA Blackwell GPUs are currently supported and tested. Support for -other GPU architectures may be added in a future release. +Only NVIDIA Blackwell GPUs (`SM100` family) are supported. The performance +results in this example were measured on GB300 NVL (`SM103`). Per-GPU memory +requirements differ per recipe, and are set by the attention layout rather +than by the GPU count: DEP16 needs 210 GB per rank and TEP8 needs 213 GB per +rank, so both require GB300-class per-GPU memory, while TEP16 needs 115 GB +per rank and is validated end-to-end on GB200 (`SM100`) at 16 GPUs. B200 +(`SM100`) is functionally supported at the kernel and module level and +covered by unit tests in CI. See Current limitations. Support for other GPU +architectures may be added in a future release. ## Prerequisites @@ -27,12 +34,14 @@ other GPU architectures may be added in a future release. for `.venv-3.12` in every command below and export `TRTLLM_VENV=/path/to/repo/.venv-.` when submitting the Slurm jobs (they default to the repository-root `.venv-3.12`). Adjust - `--cuda_architectures` to the target GPUs (`103-real` for GB300). + `--cuda_architectures` to the target GPUs (`103-real` for GB300, + `100-real` for B200). - A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded from the Hugging Face Hub (the example scripts take a local filesystem path). -- A Slurm cluster with 16 NVIDIA Blackwell GPUs and a TensorRT-LLM container +- A Slurm cluster with 16 NVIDIA Blackwell GPUs (GB300-class per-GPU memory + for DEP16 and TEP8; see Hardware support) and a TensorRT-LLM container image. The image passed as `--image` below must already provide TensorRT-LLM's runtime dependencies, that is, a release-style TensorRT-LLM container; a build or devel image without them does not work. The Slurm @@ -214,6 +223,17 @@ decoding requires the default cache manager, which cannot reuse blocks. ## Current limitations - Pipeline parallelism is not supported. +- **DEP16 and TEP8 require GB300-class per-GPU memory.** Under attention-DP + the BF16 non-expert weights are replicated on every rank (114 GB per rank), + which together with the MXFP4 routed experts at 16-way expert parallelism + (90 GB per rank) needs 210 GB per rank for DEP16; TEP8 needs 213 GB per + rank because its 8-way expert share alone is 181 GB. Neither fits `SM100` + per-GPU memory. Use TEP16, which shards the non-expert weights and needs + 115 GB per rank — validated end-to-end on GB200 (16 GPUs, 4 nodes), see the + deployment guide. Note that the FP8 weight-read path (TRTLLM-14765) does + not lift the DEP16 requirement: the conversion runs after the weights are + already resident in BF16, so it lowers the steady-state footprint but + leaves the load-time peak unchanged. - **Known performance limitation at DEP16 saturation** (attention-DP + EP16 throughput recipe): the 8K/1K serving sweep loses several percent of output throughput at concurrency ≥ 128 (up to ~15% at concurrency diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 03bcc9004801..7e1682366863 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -61,6 +61,34 @@ def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: return backend +def _validate_mla_generation_backend(backend: str, num_heads: int) -> None: + """Fail fast when `backend` can never run at this per-rank head count. + + FlashInfer's `trtllm_batch_decode_with_kv_cache_mla` rejects + `64 < num_heads_q < 128` for every batch shape, and the per-batch policy + never demotes away from an explicit `trtllm-gen` selection (the + FP8-KV-cache override or a `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen` request). + Without this check the conflict only surfaces as a FlashInfer error deep + in attention warmup. + + The bound mirrors FlashInfer's validation gate verbatim: it predicts + FlashInfer's rejection, it is not a verified support claim for the head + counts outside the range. For Kimi K3 the open side above 128 is + unreachable anyway — per-rank Q heads never exceed 96 (all heads + replicated under attention-DP, `96 / tp_size` under TEP head sharding). + """ + if backend == "trtllm-gen" and 64 < num_heads < 128: + raise ValueError( + "Kimi K3 MLA: the trtllm-gen generation backend cannot run with " + f"{num_heads} query heads per rank (trtllm-gen MLA decode rejects " + "64 < num_heads_q < 128; under attention-DP every rank keeps all " + "heads). trtllm-gen was selected explicitly — by the FP8-KV-cache " + f"override or by {_KIMI_K3_MLA_GEN_BACKEND_ENV}=trtllm-gen. Use " + "tensor-parallel head sharding (TEP) so each rank has <= 64 " + "heads, or a BF16 KV cache with the default cute-dsl backend." + ) + + def _kimi_k3_mla_decode_backend_policy( requested_backend: str, metadata: TrtllmAttentionMetadata, @@ -298,6 +326,11 @@ def __init__( # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so # only it needs K3's per-batch fallback policy; mha keeps the # default trtllm-gen selection. + # Validate here rather than in _select_mla_generation_backend: the + # per-rank head count (replicated under attention-DP, sharded under + # TEP) is only authoritative once the base MLA module has built its + # generation backend. + _validate_mla_generation_backend(self.mqa.flashinfer_mla_backend, self.mqa.num_heads) self.mqa.mla_backend_policy = partial( _kimi_k3_mla_decode_backend_policy, num_heads=self.mqa.num_heads, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4a217b4294c5..a3e8504ad21f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -106,12 +106,23 @@ l0_b200: - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py + # CuTe MTP verify kernel parity (trtllm::kda_mtp_decode). + - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py # GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the # stage's "not cpu_only" markexpr and run on the CPU-Generic stage instead). - unittest/disaggregated/test_kda_mamba_transfer.py - unittest/_torch/moe/test_kimi_k3_mlp.py - unittest/_torch/moe/test_kimi_k3_situ_moe.py + # Kimi K3 MLA backend selection + config/checkpoint plumbing (monkeypatched, + # HW-light). Collected nowhere else: the CPU-Generic stage only collects files + # carrying the literal cpu_only marker, and these carry none. + - unittest/_torch/modules/test_kimi_k3_mla_backend.py + - unittest/_torch/modeling/test_kimi_k3_config_routing.py + - unittest/_torch/modeling/test_kimi_linear_checkpoint.py + - unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py + # CPU self-test of the Kimi K3 disagg parity harness comparison logic. + - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15) # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_mla_helix.py @@ -344,6 +355,9 @@ l0_b200: # Covered by H100 pre_merge for primary HW-agnostic signal; keep B200 runtime # canary in post_merge for CUDA IPC / virtual memory / profiling paths. - unittest/_torch/misc + # Kimi K3 MoE routing shape on the TRTLLM backend (random weights, no + # checkpoint); ~13 min, too heavy for pre_merge. Mirrors l0_gb300_multi_gpus. + - unittest/_torch/moe/test_moe_backend.py -k "kimi_k3_shape" - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[mxfp8-latency] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index fb4bf93e7a9b..405a83d4db58 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -362,3 +362,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_attn_multi_gpus TIMEOUT (60) # ---- FP8 per-tensor (QDQ) moved to post-merge; block-scale/W4A8 stay in pre ---- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8" + # Kimi K3 Tier-1: truncated 4-layer (KDA + first full-attention) SA spec-dec + # logits parity on TP=4 (skips cleanly when the Kimi-K3 checkpoint is not + # staged under LLM_MODELS_ROOT). + - test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity TIMEOUT (45) diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 67ec3ea13a4b..33f32cf693c5 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -5,12 +5,17 @@ from typing import Optional import pytest +import torch +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import ( _KIMI_K3_MLA_GEN_BACKEND_ENV, _kimi_k3_mla_decode_backend_policy, _select_mla_generation_backend, + _validate_mla_generation_backend, ) +from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -53,6 +58,81 @@ def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( assert _select_mla_generation_backend(quant_config) == "trtllm-gen" +@pytest.mark.parametrize( + ("backend", "num_heads"), + [ + ("cute-dsl", 96), + ("trtllm-gen", 6), + ("trtllm-gen", 64), + ("trtllm-gen", 128), + ], +) +def test_validate_mla_generation_backend_accepts_runnable_configs( + backend: str, num_heads: int +) -> None: + _validate_mla_generation_backend(backend, num_heads) + + +@pytest.mark.parametrize("num_heads", [65, 96, 127]) +def test_validate_mla_generation_backend_rejects_trtllm_gen_mid_head_counts( + num_heads: int, +) -> None: + """trtllm-gen with 64 < H < 128 per-rank heads can never run any batch.""" + with pytest.raises(ValueError, match="query heads per rank"): + _validate_mla_generation_backend("trtllm-gen", num_heads) + + +def test_fp8_kv_cache_with_attention_dp_head_count_fails_fast( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The FP8-KV override forces trtllm-gen; with 96 per-rank heads + (attention-DP replication) construction must fail fast rather than + crash inside FlashInfer at attention warmup.""" + monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) + quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + + backend = _select_mla_generation_backend(quant_config) + with pytest.raises(ValueError, match="FP8-KV-cache"): + _validate_mla_generation_backend(backend, num_heads=96) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="KimiK3MLAAttention builds real TRTLLM attention backends", +) +def test_kimi_k3_mla_construction_fails_fast_for_fp8_kv_with_attention_dp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Constructor-level regression for the FP8-KV + attention-DP conflict: + the override selects trtllm-gen, attention-DP keeps all 96 query heads + on every rank, and `KimiK3MLAAttention.__init__` itself must raise — + before attention warmup ever runs.""" + monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) + model_config = ModelConfig( + skip_create_weights_in_init=True, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), + mapping=Mapping(world_size=4, tp_size=4, rank=0, enable_attention_dp=True), + ) + + with pytest.raises(ValueError, match="query heads per rank"): + # Kimi K3 MLA geometry (96 Q heads); small max positions keep the + # identity-RoPE table allocation negligible. + KimiK3MLAAttention( + hidden_size=7168, + num_heads=96, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rms_norm_eps=1e-6, + dtype=torch.bfloat16, + layer_idx=0, + max_position_embeddings=256, + model_config=model_config, + ) + + @pytest.mark.parametrize( ( "requested_backend",