feat: torch.compile and custom_op support

[abcdabcd987]

Follow up of #552. This PR adds torch library annotation to all FlashInfer kernels so that torch.compile can recognize the kernels. Most changes are tedious.

I manually ran subsets of pytest test cases when I made these changes, but since there are too many of them and also some of them didn't pass even before I made the change, I cannot guarantee it's all working. To run tests with torch.compile, pass FLASHINFER_TEST_TORCH_COMPILE=1 env.

# With torch.compile
FLASHINFER_TEST_TORCH_COMPILE=1 pytest -svx tests/test_norm.py

# Without torch.compile
pytest -svx tests/test_norm.py

Notable changes:

• For the prefill and decode pybind, it used to return Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] depending on return_lse. This causes trouble for torch.compile. I changed the pybind interface to accept a maybe_lse: Optional[torch.Tensor] and only return one tensor. The allocation of the lse tensor is moved to Python side. The Python API does not change.
• chain_speculative_sampling pybind: Move the allocation of accepted and emitted from C++ to Python. This is because torch.compile doesn't like returning input tensor as output tensor. The Python API does not change.

Piggyback changes:

• BatchPrefillWithRaggedKVCacheWrapper.plan: Bugfix qo_indptr not on CPU
• merge_state: Fix typo in docs
• Change run_return_lse(...) to run(..., return_lse=True) because torch.compile does not recognize functools.partial.
• In tests, change flashinfer.xxx() to flashinfer.<module>.xxx() so that the monkeypatch works.

Unsupported for torch.compile:

• flashinfer.quantization.segment_packbits: Because it's data dependent.

Untouched:

• sparse.py: Tests didn't pass beforehand, so I skiped this. Also, it doesn't seem like need custom_op annotations, as it does not have CUDA kernels.

Failed test cases:

• batch_decode non contiguous kv: test_batch_decode_with_paged_kv_cache[False-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-128-4-4-1-54-12]

[yzh119]

LGTM, thanks for the huge improvement, I have left some tiny suggestions.