From 4ac41bbd3773e085d9c88faf67b0aad8e9ea81d5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 7 Jul 2026 19:04:18 -0700 Subject: [PATCH 1/3] [None][fix] Use per-layer page-index scale for FlashInfer KV page tables When max_seq_len is at most the sliding window size, every attention window clamps to max_seq_len, so all layers land in one KV cache pool. Layers with different geometry (Gemma4: sliding head_dim 256, global head_dim 512) then share that pool while having different page-index scales, but page indices were converted with the single pool-level scale and shared across all layers. Global-attention layers received page ids up to 4x past their buffer, and append_paged_kv_cache crashed with CUDA_ERROR_ILLEGAL_ADDRESS during the warmup prefill on every Gemma4 checkpoint whenever max_seq_len <= 512. Default (large) max_seq_len was unaffected: pools then split per window and each pool is homogeneous. - kv_cache_manager_v2: convert per-layer page-index requests with that layer own scale (the V2 core documents that computed page indices may only be shared between buffers with equal scale). - FlashInfer metadata: share one page-index list only among layers with the same pool and scale; engage the per-pool machinery when windows differ (VSWA) or scales differ. - Add a Gemma4 dummy-weight regression test at max_seq_len 256 and 512 (real E2B geometry), registered in the B200 pre-merge list. Signed-off-by: tianruih --- .../_torch/attention_backend/flashinfer.py | 29 +++++++++---- .../_torch/pyexecutor/kv_cache_manager_v2.py | 13 +++++- .../test_lists/test-db/l0_b200.yml | 2 + .../_torch/modeling/test_gemma4_e2e_dummy.py | 42 +++++++++++++++++++ 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 4633850204b6..75dd2c99587a 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -550,17 +550,28 @@ def _post_init_with_buffers(self, buffers) -> None: capture_graph=capture_graph, ) - # Detect VSWA: check if the manager has multiple pools. - # Guard on layer_to_pool_mapping_dict which is V2-specific — V1 - # managers also expose is_vswa but lack the per-pool infrastructure. - if (getattr(self.kv_cache_manager, 'is_vswa', False) and hasattr( - self.kv_cache_manager, 'layer_to_pool_mapping_dict')): - mgr = self.kv_cache_manager - self._vswa_layer_to_pool = {} - self._vswa_pool_to_rep_layer: Dict[int, int] = {} + # Layers may share one page-index list only when they are in the + # same pool AND have the same page-index scale: VSWA splits pools, + # and per-layer geometry (e.g. Gemma4 sliding/global head_dim) + # splits scales even when all windows collapse to max_seq_len and + # is_vswa is False. Guarded on V2-specific attributes (V1 managers + # lack the per-pool infrastructure). + mgr = self.kv_cache_manager + _get_scale = getattr(mgr, 'get_layer_page_index_scale', None) + _layer_space: Dict[int, int] = {} + if hasattr(mgr, 'layer_to_pool_mapping_dict') and _get_scale: + _space_ids = {} for layer_idx in getattr(mgr, 'layer_offsets', {}): layer_offset = mgr.layer_offsets[layer_idx] - pool_id = mgr.layer_to_pool_mapping_dict[layer_offset] + key = (mgr.layer_to_pool_mapping_dict[layer_offset], + _get_scale(layer_idx)) + _space_ids.setdefault(key, len(_space_ids)) + _layer_space[layer_idx] = _space_ids[key] + if _layer_space and (getattr(mgr, 'is_vswa', False) + or len(set(_layer_space.values())) > 1): + self._vswa_layer_to_pool = {} + self._vswa_pool_to_rep_layer: Dict[int, int] = {} + for layer_idx, pool_id in _layer_space.items(): self._vswa_layer_to_pool[layer_idx] = pool_id if pool_id not in self._vswa_pool_to_rep_layer: self._vswa_pool_to_rep_layer[pool_id] = layer_idx diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index dbf0d983c66e..31fed861250d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2775,6 +2775,12 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): else: self.index_mapper.remove_sequence(request.py_request_id) + def get_layer_page_index_scale(self, layer_idx: int) -> int: + """Page-index scale of this layer's KV buffer. Layers in one pool can + have different scales (e.g. different head_dim), so per-layer callers + must not use the pool-level scale.""" + return int(self.impl.get_page_index_scale(self.layer_offsets[layer_idx], Role.KEY)) + def get_batch_cache_indices( self, request_ids: List[int], @@ -2783,13 +2789,16 @@ def get_batch_cache_indices( ) -> List[List[int]]: if layer_idx is None: pool_id = 0 + index_scale = None else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] + index_scale = self.get_layer_page_index_scale(layer_idx) return self._get_batch_cache_indices_by_pool_id( request_ids, pool_id=pool_id, is_kv_aggregate=True, num_blocks_per_seq=num_blocks_per_seq, + index_scale=index_scale, ) def _get_batch_cache_indices_by_pool_id( @@ -2799,6 +2808,7 @@ def _get_batch_cache_indices_by_pool_id( pool_id: int = 0, is_kv_aggregate: bool = True, num_blocks_per_seq: Optional[Sequence[int]] = None, + index_scale: Optional[int] = None, ) -> List[List[int]]: if is_kv_aggregate: # Div by kv_factor to index kv cache with size @@ -2807,7 +2817,8 @@ def _get_batch_cache_indices_by_pool_id( else: div_factor = 1 - index_scale = int(self.index_scales[pool_id]) + if index_scale is None: + index_scale = int(self.index_scales[pool_id]) res = [] for req_idx, req_id in enumerate(request_ids): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index cf51c3c05a6e..c9428acb095e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -169,6 +169,8 @@ l0_b200: - unittest/_torch/modeling/test_gemma4_multimodal.py - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_26b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e2b_dummy + - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e2b_dummy_small_max_seq_len[256] + - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e2b_dummy_small_max_seq_len[512] - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e4b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_multimodal_26b_dummy diff --git a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py index a728b892d0a4..16e1bd9b7b43 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py +++ b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py @@ -226,6 +226,48 @@ def test_e2e_text_31b_dummy(): shutil.rmtree(dummy_dir, ignore_errors=True) +@requires_gemma4_transformers +@pytest.mark.skipif(not _model_available("E2B"), reason="gemma-4-E2B-it not found") +@pytest.mark.parametrize("max_seq_len", [256, 512]) +def test_e2e_text_e2b_dummy_small_max_seq_len(max_seq_len): + """E2E with max_seq_len <= sliding_window (256 and 512). + + Regression test: when max_seq_len is at most the sliding window size, all + attention windows clamp to max_seq_len, the KV cache manager reports a + single window (is_vswa False), and the FlashInfer metadata used to skip + the per-pool page-index mapping — while the pools still differ by + head_dim. The shared page-index list then sends out-of-range page ids to + the smaller (global head_dim) pool and append_paged_kv_cache crashes with + an illegal memory access during the warmup prefill. + """ + from tensorrt_llm.llmapi import LLM, SamplingParams + + # Real E2B geometry with dummy weights: the sliding (head_dim 256) and + # global (head_dim 512) layers must keep different page-index scales, + # and sliding_window must stay 512 so both max_seq_len values are at + # most the window. Shrinking the config changes both and hides the bug. + dummy_dir = _make_dummy_config_dir( + MODEL_PATHS["E2B"], + dummy_head_dim=256, + dummy_global_head_dim=512, + shrink_hidden=False, + ) + try: + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + max_seq_len=max_seq_len, + ) + with llm: + output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) + assert len(output) == 1 + assert len(output[0].outputs[0].token_ids) > 0 + finally: + shutil.rmtree(dummy_dir, ignore_errors=True) + + @requires_gemma4_transformers @pytest.mark.skipif(not _model_available("E4B"), reason="gemma-4-E4B-it not found") def test_e2e_text_e4b_dummy(): From c8f0d2753c93cab561a2e817efe62653a8e37ea7 Mon Sep 17 00:00:00 2001 From: Tianrui Hu Date: Mon, 13 Jul 2026 22:33:51 -0700 Subject: [PATCH 2/3] [None][fix] Gemma4 IMA test/backend: address review feedback - flashinfer: drop leading underscores from the page-index-scale locals - gemma4 e2e test: fail (not skip) when a checkpoint is missing, since CI runs these with LLM_MODELS_ROOT set; a missing artifact is a setup error - gemma4 e2e test: hoist the LLM/SamplingParams import to module level Signed-off-by: Tianrui Hu --- .../_torch/attention_backend/flashinfer.py | 20 +++++----- .../_torch/modeling/test_gemma4_e2e_dummy.py | 39 +++++++------------ 2 files changed, 25 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 7b281cdfd3ed..48714e841c4c 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -670,21 +670,21 @@ def _post_init_with_buffers(self, buffers) -> None: # is_vswa is False. Guarded on V2-specific attributes (V1 managers # lack the per-pool infrastructure). mgr = self.kv_cache_manager - _get_scale = getattr(mgr, 'get_layer_page_index_scale', None) - _layer_space: Dict[int, int] = {} - if hasattr(mgr, 'layer_to_pool_mapping_dict') and _get_scale: - _space_ids = {} + get_scale = getattr(mgr, 'get_layer_page_index_scale', None) + layer_space: Dict[int, int] = {} + if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale: + space_ids = {} for layer_idx in getattr(mgr, 'layer_offsets', {}): layer_offset = mgr.layer_offsets[layer_idx] key = (mgr.layer_to_pool_mapping_dict[layer_offset], - _get_scale(layer_idx)) - _space_ids.setdefault(key, len(_space_ids)) - _layer_space[layer_idx] = _space_ids[key] - if _layer_space and (getattr(mgr, 'is_vswa', False) - or len(set(_layer_space.values())) > 1): + get_scale(layer_idx)) + space_ids.setdefault(key, len(space_ids)) + layer_space[layer_idx] = space_ids[key] + if layer_space and (getattr(mgr, 'is_vswa', False) + or len(set(layer_space.values())) > 1): self._vswa_layer_to_pool = {} self._vswa_pool_to_rep_layer: Dict[int, int] = {} - for layer_idx, pool_id in _layer_space.items(): + for layer_idx, pool_id in layer_space.items(): self._vswa_layer_to_pool[layer_idx] = pool_id if pool_id not in self._vswa_pool_to_rep_layer: self._vswa_pool_to_rep_layer[pool_id] = layer_idx diff --git a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py index 16e1bd9b7b43..427a9a7680ea 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py +++ b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py @@ -44,6 +44,10 @@ pytest.skip("LLM_MODELS_ROOT not set", allow_module_level=True) _GEMMA4_MODELS = os.path.join(_LLM_MODELS_ROOT, "gemma4") +# Imported after the module-level skip guard so that collecting this module on +# a machine without LLM_MODELS_ROOT does not pull in the runtime import. +from tensorrt_llm.llmapi import LLM, SamplingParams # noqa: E402 + # Real model paths — used for tokenizer + base config MODEL_PATHS = { "26B": os.path.join(_GEMMA4_MODELS, "gemma-4-26B-A4B-it"), @@ -53,11 +57,6 @@ } -def _model_available(name: str) -> bool: - path = MODEL_PATHS.get(name, "") - return os.path.isfile(os.path.join(path, "config.json")) - - def _make_dummy_config_dir( model_path: str, dummy_head_dim: int = 128, @@ -73,10 +72,20 @@ def _make_dummy_config_dir( Returns path to the temp directory. """ + # Fail loudly rather than silently skipping when the real checkpoint is + # missing: these tests are registered in CI with LLM_MODELS_ROOT set, so an + # absent artifact is a setup error, not a reason to report a passing skip. + config_path = os.path.join(model_path, "config.json") + if not os.path.isfile(config_path): + raise FileNotFoundError( + f"Gemma4 test checkpoint not found: {config_path}. These E2E tests " + f"require the real gemma4 models under $LLM_MODELS_ROOT/gemma4." + ) + tmp_dir = tempfile.mkdtemp(prefix="gemma4_dummy_") # Load and patch config - with open(os.path.join(model_path, "config.json")) as f: + with open(config_path) as f: config = json.load(f) tc = config.get("text_config", config) @@ -176,11 +185,8 @@ def _make_dummy_config_dir( @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("26B"), reason="gemma-4-26B-A4B-it not found") def test_e2e_text_26b_dummy(): """E2E text generation for 26B-A4B (MoE + K=V + softcap + hybrid attn).""" - from tensorrt_llm.llmapi import LLM, SamplingParams - dummy_dir = _make_dummy_config_dir(MODEL_PATHS["26B"]) try: llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") @@ -193,11 +199,8 @@ def test_e2e_text_26b_dummy(): @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("E2B"), reason="gemma-4-E2B-it not found") def test_e2e_text_e2b_dummy(): """E2E text generation for E2B (KV sharing + PLE + double-wide MLP).""" - from tensorrt_llm.llmapi import LLM, SamplingParams - dummy_dir = _make_dummy_config_dir(MODEL_PATHS["E2B"]) try: llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") @@ -210,11 +213,8 @@ def test_e2e_text_e2b_dummy(): @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("31B"), reason="gemma-4-31B-it not found") def test_e2e_text_31b_dummy(): """E2E text generation for 31B (K=V + hybrid attn + softcap).""" - from tensorrt_llm.llmapi import LLM, SamplingParams - dummy_dir = _make_dummy_config_dir(MODEL_PATHS["31B"]) try: llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") @@ -227,7 +227,6 @@ def test_e2e_text_31b_dummy(): @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("E2B"), reason="gemma-4-E2B-it not found") @pytest.mark.parametrize("max_seq_len", [256, 512]) def test_e2e_text_e2b_dummy_small_max_seq_len(max_seq_len): """E2E with max_seq_len <= sliding_window (256 and 512). @@ -240,8 +239,6 @@ def test_e2e_text_e2b_dummy_small_max_seq_len(max_seq_len): the smaller (global head_dim) pool and append_paged_kv_cache crashes with an illegal memory access during the warmup prefill. """ - from tensorrt_llm.llmapi import LLM, SamplingParams - # Real E2B geometry with dummy weights: the sliding (head_dim 256) and # global (head_dim 512) layers must keep different page-index scales, # and sliding_window must stay 512 so both max_seq_len values are at @@ -269,11 +266,8 @@ def test_e2e_text_e2b_dummy_small_max_seq_len(max_seq_len): @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("E4B"), reason="gemma-4-E4B-it not found") def test_e2e_text_e4b_dummy(): """E2E text generation for E4B (KV sharing + hybrid attn).""" - from tensorrt_llm.llmapi import LLM, SamplingParams - dummy_dir = _make_dummy_config_dir(MODEL_PATHS["E4B"]) try: llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") @@ -291,14 +285,11 @@ def test_e2e_text_e4b_dummy(): @requires_gemma4_transformers -@pytest.mark.skipif(not _model_available("26B"), reason="gemma-4-26B-A4B-it not found") def test_e2e_multimodal_26b_dummy(): """E2E multimodal: image → vision tower → embedder → LLM → output.""" import numpy as np from PIL import Image - from tensorrt_llm.llmapi import LLM, SamplingParams - dummy_dir = _make_dummy_config_dir(MODEL_PATHS["26B"]) try: # Format prompt with image placeholder via tokenizer chat template From f80bb6420c9d669609cc8b4bd0a5f98ed04cccdf Mon Sep 17 00:00:00 2001 From: Tianrui Hu Date: Tue, 14 Jul 2026 03:10:03 -0700 Subject: [PATCH 3/3] [None][fix] Gemma4 IMA: per-layer scale in flat block-table path + correct model root - kv_cache_manager_v2: get_batch_cache_indices_flat now uses the per-layer page-index scale when layer_idx is given (matches get_batch_cache_indices). prepare() builds VSWA pool block tables through this flat helper, so the earlier fix to the non-flat path alone left the trtllm-gen decode path still scaling shared-pool layers with the pool-level scale -> out-of-range page ids -> CUDA_ERROR_ILLEGAL_ADDRESS during warmup prefill. - test_gemma4_e2e_dummy: use the canonical model root subdir 'gemma' (per tests/test_common/llm_data.py) instead of 'gemma4', so the E2E tests actually find the checkpoints instead of silently skipping. Signed-off-by: Tianrui Hu --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 8 ++- .../_torch/modeling/test_gemma4_e2e_dummy.py | 55 ++++++++++++++++--- 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index ffc24cc18a79..16f4e08a7e10 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2882,10 +2882,14 @@ def get_batch_cache_indices_flat( """ if layer_idx is None: pool_id = 0 + scale = self._index_scale_ints[pool_id] else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] - - scale = self._index_scale_ints[pool_id] + # Layers sharing a pool can still require different page-index + # scales (e.g. Gemma4 sliding/global head_dim). Use the per-layer + # scale so this flat block table matches get_batch_cache_indices() + # and never feeds out-of-range page ids to FlashInfer. + scale = self.get_layer_page_index_scale(layer_idx) div_factor = self.kv_factor out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) diff --git a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py index 427a9a7680ea..5637879e926a 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py +++ b/tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py @@ -42,11 +42,19 @@ _LLM_MODELS_ROOT = os.environ.get("LLM_MODELS_ROOT") if _LLM_MODELS_ROOT is None: pytest.skip("LLM_MODELS_ROOT not set", allow_module_level=True) -_GEMMA4_MODELS = os.path.join(_LLM_MODELS_ROOT, "gemma4") +# Canonical model root subdir is "gemma" (see tests/test_common/llm_data.py: +# "google/gemma-4-E2B-it" -> "gemma/gemma-4-E2B-it"), not "gemma4". +_GEMMA4_MODELS = os.path.join(_LLM_MODELS_ROOT, "gemma") # Imported after the module-level skip guard so that collecting this module on # a machine without LLM_MODELS_ROOT does not pull in the runtime import. -from tensorrt_llm.llmapi import LLM, SamplingParams # noqa: E402 +from tensorrt_llm.llmapi import LLM, KvCacheConfig, SamplingParams # noqa: E402 + +# These dummy models are tiny, but the default KV-cache fraction sizes the pool +# to most of the (very large B200) device memory, leaving nothing for the other +# executor components -> OOM at executor creation. Cap it so the whole pipeline +# fits regardless of card size. +_KV_CACHE_CONFIG = KvCacheConfig(free_gpu_memory_fraction=0.5) # Real model paths — used for tokenizer + base config MODEL_PATHS = { @@ -79,7 +87,7 @@ def _make_dummy_config_dir( if not os.path.isfile(config_path): raise FileNotFoundError( f"Gemma4 test checkpoint not found: {config_path}. These E2E tests " - f"require the real gemma4 models under $LLM_MODELS_ROOT/gemma4." + f"require the real gemma-4 models under $LLM_MODELS_ROOT/gemma." ) tmp_dir = tempfile.mkdtemp(prefix="gemma4_dummy_") @@ -189,7 +197,13 @@ def test_e2e_text_26b_dummy(): """E2E text generation for 26B-A4B (MoE + K=V + softcap + hybrid attn).""" dummy_dir = _make_dummy_config_dir(MODEL_PATHS["26B"]) try: - llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + kv_cache_config=_KV_CACHE_CONFIG, + ) with llm: output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) assert len(output) == 1 @@ -203,7 +217,13 @@ def test_e2e_text_e2b_dummy(): """E2E text generation for E2B (KV sharing + PLE + double-wide MLP).""" dummy_dir = _make_dummy_config_dir(MODEL_PATHS["E2B"]) try: - llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + kv_cache_config=_KV_CACHE_CONFIG, + ) with llm: output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) assert len(output) == 1 @@ -217,7 +237,13 @@ def test_e2e_text_31b_dummy(): """E2E text generation for 31B (K=V + hybrid attn + softcap).""" dummy_dir = _make_dummy_config_dir(MODEL_PATHS["31B"]) try: - llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + kv_cache_config=_KV_CACHE_CONFIG, + ) with llm: output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) assert len(output) == 1 @@ -256,6 +282,7 @@ def test_e2e_text_e2b_dummy_small_max_seq_len(max_seq_len): attn_backend="FLASHINFER", dtype="bfloat16", max_seq_len=max_seq_len, + kv_cache_config=_KV_CACHE_CONFIG, ) with llm: output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) @@ -270,7 +297,13 @@ def test_e2e_text_e4b_dummy(): """E2E text generation for E4B (KV sharing + hybrid attn).""" dummy_dir = _make_dummy_config_dir(MODEL_PATHS["E4B"]) try: - llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + kv_cache_config=_KV_CACHE_CONFIG, + ) with llm: output = llm.generate(["Hello"], SamplingParams(max_tokens=4)) assert len(output) == 1 @@ -310,7 +343,13 @@ def test_e2e_multimodal_26b_dummy(): tokenize=False, ) - llm = LLM(dummy_dir, load_format="dummy", attn_backend="FLASHINFER", dtype="bfloat16") + llm = LLM( + dummy_dir, + load_format="dummy", + attn_backend="FLASHINFER", + dtype="bfloat16", + kv_cache_config=_KV_CACHE_CONFIG, + ) with llm: img = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)) prompt = {