FAILED tests/attention/test_trtllm_gen_attention.py::test_trtllm_batch_decode_spec[True-True-110-False-None-fp8-fp8-nvfp4-127-256-16-32-2-8-256-NHD-trtllm-gen] - torch.AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
========= 73964 failed, 12 passed, 100192 skipped in 218.63s (0:03:38) =========
❌ FAILED: tests/attention/test_trtllm_gen_attention.py

/workspace/flashinfer$ pytest tests/attention/test_trtllm_gen_attention.py -x
===================================================== test session starts =====================================================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /workspace/flashinfer
configfile: pytest.ini
collected 174168 items

tests/attention/test_trtllm_gen_attention.py ............F

========================================================== FAILURES ===========================================================
_____________ test_trtllm_batch_prefill[True-False-False-128-2047-511-True-None-bf16-bf16-bf16--1-256-16-4-8-HND] _____________

kv_layout = 'HND', batch_size = 256, page_size = 16, num_kv_heads = 4, head_grp_size = 8, window_left = -1, q_dtype = 'bf16'
o_dtype = 'bf16', kv_dtype = 'bf16', enable_pdl = None, enable_sink = True, max_q_len = 511, max_kv_len = 2047, head_dim = 128
non_contiguous_query = False, skips_softmax = False, uses_shared_paged_kv_idx = True

    @pytest.mark.parametrize("kv_layout", ["HND", "NHD"])
    @pytest.mark.parametrize(
        "batch_size,page_size,num_kv_heads,head_grp_size",
        [
            (4, 16, 2, 1),
            (4, 32, 4, 5),
            (4, 64, 4, 8),
            (128, 16, 2, 5),
            (128, 32, 4, 1),
            (128, 64, 2, 8),
            (256, 16, 4, 8),
            (256, 32, 2, 8),
            (256, 64, 4, 1),
            (256, 64, 4, 5),
        ],
    )
    @pytest.mark.parametrize("window_left", [-1])  # todo(Siyuan): add 127 window_left
    @pytest.mark.parametrize(
        "q_dtype,kv_dtype,o_dtype",
        [
            ("bf16", "bf16", "bf16"),
            ("fp16", "fp16", "fp16"),
            ("fp8", "fp8", "bf16"),
            ("fp8", "fp8", "fp16"),
            ("fp8", "fp8", "fp8"),
            ("fp8", "fp8", "nvfp4"),
            ("fp8", "nvfp4", "fp8"),
        ],
    )
    @pytest.mark.parametrize("enable_pdl", [None])
    @pytest.mark.parametrize("enable_sink", [True, False])
    @pytest.mark.parametrize("max_q_len", [511])
    @pytest.mark.parametrize("max_kv_len", [2047])
    @pytest.mark.parametrize("head_dim", [128, 256])
    @pytest.mark.parametrize("non_contiguous_query", [False, True])
    @pytest.mark.parametrize("skips_softmax", [False, True])
    @pytest.mark.parametrize("uses_shared_paged_kv_idx", [True, False])
    def test_trtllm_batch_prefill(
        kv_layout: str,
        batch_size: int,
        page_size: int,
        num_kv_heads: int,
        head_grp_size: int,
        window_left: int,
        q_dtype: str,
        o_dtype: str,
        kv_dtype: str,
        enable_pdl: bool,
        enable_sink: bool,
        max_q_len: int,
        max_kv_len: int,
        head_dim: int,
        non_contiguous_query: bool,
        skips_softmax: bool,
        uses_shared_paged_kv_idx: bool,
    ):
>       _test_trtllm_batch_prefill(
            kv_layout,
            batch_size,
            page_size,
            num_kv_heads,
            head_grp_size,
            window_left,
            q_dtype,
            o_dtype,
            kv_dtype,
            enable_pdl,
            enable_sink,
            max_q_len,
            max_kv_len,
            kv_dtype in ("fp8", "nvfp4"),
            head_dim,
            non_contiguous_query=non_contiguous_query,
            skips_softmax=skips_softmax,
            uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
        )

tests/attention/test_trtllm_gen_attention.py:852:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

kv_layout = 'HND', batch_size = 256, page_size = 16, num_kv_heads = 4, head_grp_size = 8, window_left = -1, q_dtype = 'bf16'
o_dtype = 'bf16', kv_dtype = 'bf16', enable_pdl = None, enable_sink = True, max_q_len = 511, max_kv_len = 2047
device_scale = False, head_dim = 128, non_contiguous_query = False, skips_softmax = False, uses_shared_paged_kv_idx = True

    def _test_trtllm_batch_prefill(
        kv_layout: str,
        batch_size: int,
        page_size: int,
        num_kv_heads: int,
        head_grp_size: int,
        window_left: int,
        q_dtype: str,
        o_dtype: str,
        kv_dtype: str,
        enable_pdl: bool,
        enable_sink: bool,
        max_q_len: int,
        max_kv_len: int,
        device_scale: float,
        head_dim: int,
        non_contiguous_query: bool = False,
        skips_softmax: bool = False,
        uses_shared_paged_kv_idx: bool = True,
    ):
        compute_capability = get_compute_capability(torch.device(device="cuda"))
        if compute_capability[0] != 10:
            pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")

        if skips_softmax and q_dtype != kv_dtype:
            pytest.skip(
                "skips_softmax does not currently support Q and Kv types being different"
            )

        # NVFP4 KV cache constraints
        if kv_dtype == "nvfp4":
            if q_dtype != "fp8":
                pytest.skip("NVFP4 KV cache requires FP8 query")
            if o_dtype != "fp8":
                pytest.skip("NVFP4 KV cache only supports FP8 output")

        # Set up test parameters
        torch.manual_seed(0)

        # Generate random sequence lengths
        num_qo_heads = num_kv_heads * head_grp_size
        q_lens, _, seq_lens = generate_seq_lens_prefill(batch_size, max_q_len, max_kv_len)

        # Create query tensor and related data
        q, q_scale, ref_q = create_query_tensor(q_lens, num_qo_heads, head_dim, q_dtype)
        q_indptr = generate_cumsum_lens(q_lens)

        # Create KV cache and related data
        kv_cache, k_scale, v_scale, ref_kv_cache, kv_cache_sf = create_kv_cache(
            batch_size,
            seq_lens,
            page_size,
            num_kv_heads,
            head_dim,
            kv_dtype,
            "bf16" if q_dtype == "fp8" or kv_dtype == "nvfp4" else q_dtype,
            kv_layout,
        )
        page_table, all_page_ids, page_per_seq = create_page_table(
            batch_size, seq_lens, page_size
        )
        kv_indptr = generate_cumsum_lens(page_per_seq)
        kv_last_page_len = get_last_page_len(seq_lens, page_size)

        kv_cache_kernel, page_table_kernel, kv_cache_sf_kernel = (
            prepare_paged_kv_for_kernel(
                kv_cache, page_table, uses_shared_paged_kv_idx, kv_cache_sf
            )
        )

        workspace_buffer, workspace_buffer_ref = create_workspace_buffers(GPU_DEVICE)

        # Create output tensor and related data
        create_out_tensor = flip_coin(
            batch_size, page_size, num_kv_heads, head_grp_size, o_dtype
        )
        can_infer_type = q.dtype == DTYPE_MAP[o_dtype] or create_out_tensor
        create_out_dtype = not can_infer_type or flip_coin(
            batch_size, page_size, num_kv_heads, head_grp_size, o_dtype, q_dtype
        )
        out, out_dtype, o_scale, o_sf_scale, o_sf_vec_size = create_output(
            q, o_dtype, create_out_tensor, create_out_dtype
        )

        sm_scale = float(1.0 / (head_dim**0.5))
        lse_ref = None

        # Build reference output
        plan_params = {
            "qo_indptr": q_indptr,
            "paged_kv_indptr": kv_indptr,
            "paged_kv_indices": all_page_ids,
            "paged_kv_last_page_len": kv_last_page_len.to(GPU_DEVICE),
            "num_qo_heads": num_qo_heads,
            "num_kv_heads": num_kv_heads,
            "head_dim_qk": head_dim,
            "page_size": page_size,
            "causal": True,
            "pos_encoding_mode": "NONE",
            "logits_soft_cap": 0.0,
            "q_data_type": ref_q.dtype,
            "kv_data_type": ref_kv_cache.dtype,
            "window_left": window_left,
        }
        if not enable_sink:
            wrapper_ref = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
                workspace_buffer_ref, kv_layout
            )
            wrapper_ref.plan(**plan_params)
            output_ref, lse_ref = wrapper_ref.run(ref_q, ref_kv_cache, return_lse=True)
        else:
            # Construct flat K/V via helper
            k_flat, v_flat, kv_indptr_tokens = flatten_paged_kv(
                ref_kv_cache,
                page_table,
                seq_lens.to(GPU_DEVICE),
                page_size,
                kv_last_page_len,
                kv_layout,
            )
            sink = torch.rand(num_qo_heads, device=GPU_DEVICE, dtype=torch.float32) * 5
            output_ref = sink_attention_unified(
                ref_q,
                k_flat,
                v_flat,
                sink,
                window_left,
                True,
                sm_scale,
                mode="varlen",
                batch_size=batch_size,
                qo_indptr=q_indptr,
                kv_indptr=kv_indptr_tokens,
            )

        # Run trtllm-gen function call
        bmm1_scale = q_scale * k_scale * sm_scale
        bmm2_scale = v_scale / o_scale
        if isinstance(bmm1_scale, torch.Tensor) and not device_scale:
            bmm1_scale = bmm1_scale.item()
        elif not isinstance(bmm1_scale, torch.Tensor) and device_scale:
            bmm1_scale = torch.tensor(bmm1_scale, device=GPU_DEVICE, dtype=torch.float32)
        if isinstance(bmm2_scale, torch.Tensor) and not device_scale:
            bmm2_scale = bmm2_scale.item()
        elif not isinstance(bmm2_scale, torch.Tensor) and device_scale:
            bmm2_scale = torch.tensor(bmm2_scale, device=GPU_DEVICE, dtype=torch.float32)

        # Optionally make query non-contiguous for testing stride support
        if non_contiguous_query:
            q_input = make_query_non_contiguous(q, num_qo_heads, head_dim)
        else:
            q_input = q.contiguous()

        # Using a tiny threshold should give the same result as normal attention.
        skip_softmax_threshold_scale_factor = 1e-30 if skips_softmax else None
        softmax_end_bytes = trtllm_gen_workspace_softmax_end_bytes_context(
            workspace_buffer,
            num_qo_heads=q_input.size(1),
            sum_seq_q=q_input.size(0),
        )
        workspace_check_end_bytes = min(
            softmax_end_bytes + TRTLLM_GEN_WORKSPACE_CHECK_BYTES, workspace_buffer.numel()
        )
        workspace_buffer[softmax_end_bytes:workspace_check_end_bytes].zero_()
        provided_lse = torch.empty(
            (q_input.size(0), q_input.size(1)),
            device=GPU_DEVICE,
            dtype=torch.float32,
        )

        output, lse = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
            q_input,
            kv_cache_kernel,
            workspace_buffer,
            page_table_kernel,
            seq_lens.to(GPU_DEVICE),
            torch.max(q_lens).item(),
            torch.max(seq_lens).item(),
            bmm1_scale,  # bmm1_scale
            bmm2_scale,  # bmm2_scale
            batch_size,
            q_indptr,
            kv_indptr,
            window_left,  # window_left
            out=out,
            out_dtype=out_dtype,
            o_sf_scale=o_sf_scale,
            o_sf_vec_size=o_sf_vec_size,
            kv_layout=kv_layout,
            enable_pdl=enable_pdl,
            sinks=(sink if enable_sink else None),
            kv_cache_sf=kv_cache_sf_kernel,
            skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
            uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
            lse=provided_lse,
            return_lse=True,
        )
        assert (
            workspace_buffer[softmax_end_bytes:workspace_check_end_bytes].cpu().numpy() == 0
        ).all()

        if o_dtype == "nvfp4":
            output, output_ref = unpack_compare_nvfp4(
                output, output_ref, o_sf_scale, o_sf_vec_size
            )
            assert o_scale == 1.0
            rtol, atol = 4e-1, 1e0
        elif q_dtype == "fp8" and o_dtype == "fp8":
            rtol, atol = 5e-2, 7e-2
        elif q_dtype == "fp8" and o_dtype in ["bf16", "fp16"]:
            rtol, atol = 4e-2, 6e-2
        else:
            rtol, atol = 1e-2, 1e-2

        # NVFP4 KV cache has significant quantization error, especially with
        # outlier channels that create large per-block dynamic range.
        if kv_dtype == "nvfp4":
            rtol, atol = 5e-1, 5e-1

        # NVFP4 KV cache has higher mismatch rate due to 4-bit quantization noise,
        # especially with outlier channels that stress per-block scaling.
        allowed_mismatch_rate = 0.10 if kv_dtype == "nvfp4" else 1e-7
        # Calculate max allowed mismatched elements based on tensor size
        total_elements = (output.float() * o_scale).numel()
        max_mismatched_elements = int(allowed_mismatch_rate * total_elements)

        # convert to float32 for fp8 is not supported by assert_close
        assert_close_with_mismatch_tolerance(
            output.float() * o_scale,
            output_ref.float(),
            rtol=rtol,
            atol=atol,
            max_mismatched_elements=max_mismatched_elements,
        )

        # NVFP4 KV cache: use cosine similarity to catch block-scale mismatches
        # (e.g. wrong swizzling) that element-wise tolerances miss.
        if kv_dtype == "nvfp4":
            cos = torch.nn.functional.cosine_similarity(
                (output.float() * o_scale).reshape(-1),
                output_ref.float().reshape(-1),
                dim=0,
            )
            assert cos.item() > 0.86, (
                f"NVFP4 KV cache attention: cosine similarity {cos:.4f} < 0.86. "
                f"Block scale factors may be mismatched to FP4 data blocks."
            )

        expected_lse_shape = (q_input.size(0), q_input.size(1))
        assert lse is provided_lse
        assert lse.shape == expected_lse_shape
        assert lse.dtype == torch.float32
        assert torch.isfinite(lse).all()
        if lse_ref is not None:
            lse_rtol, lse_atol = get_lse_test_tolerances(q_dtype, kv_dtype)
            torch.testing.assert_close(lse, lse_ref, rtol=lse_rtol, atol=lse_atol)

        if (
            o_dtype != "nvfp4" and kv_dtype != "nvfp4" and uses_shared_paged_kv_idx
        ):  # wrapper api does not support fp4 output/kv or separate KV page indices yet.
            # test wrapper with trtllm-gen backend
            wrapper_trtllm_gen = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(
                workspace_buffer, kv_layout, backend="trtllm-gen"
            )
            plan_params["q_data_type"] = q.dtype
            plan_params["kv_data_type"] = kv_cache.dtype
            plan_params["o_data_type"] = DTYPE_MAP[o_dtype]
            wrapper_trtllm_gen.plan(**plan_params)
            workspace_buffer[softmax_end_bytes:workspace_check_end_bytes].zero_()
            output_wrapper = wrapper_trtllm_gen.run(
                q_input,
                kv_cache,
                q_scale=q_scale,
                k_scale=k_scale,
                v_scale=v_scale / o_scale,
                enable_pdl=enable_pdl,
                sinks=(sink if enable_sink else None),
            )
            # v_scale, o_scale in wrapper is emulated by multiplying output by v_scale instead of fused into kernel.
            if v_scale == o_scale == 1.0:
>               assert (output_wrapper == output).all()
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E               torch.AcceleratorError: CUDA error: an illegal memory access was encountered
E               Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
E               CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
E               For debugging consider passing CUDA_LAUNCH_BLOCKING=1
E               Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

tests/attention/test_trtllm_gen_attention.py:785: AcceleratorError
=================================================== short test summary info ===================================================
FAILED tests/attention/test_trtllm_gen_attention.py::test_trtllm_batch_prefill[True-False-False-128-2047-511-True-None-bf16-bf16-bf16--1-256-16-4-8-HND] - torch.AcceleratorError: CUDA error: an illegal memory access was encountered
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
================================================ 1 failed, 12 passed in 28.75s

FAILED tests/attention/test_trtllm_gen_attention.py::test_trtllm_batch_prefill[True-False-False-128-2047-511-True-None-bf16-bf16-bf16--1-256-16-4-8-HND] - torch.AcceleratorError: CUDA error: an illegal memory access was encountered