Add --attention_backend CLI argument (fixes minor note in #46)

[RollerCoaster1899]

Summary

Replaces hardcoded ATTENTION_BACKEND = "fa3", PAGE_SIZE, and MEM_FRACTION_STATIC constants with configurable CLI arguments in infer.py.

Problem

infer.py hardcoded ATTENTION_BACKEND = "fa3", which is FlashAttention-3 (Hopper-only, sm_90). Users with Ada GPUs (e.g. RTX 4090, sm_89) had to manually edit the source to use flashinfer or triton.

Changes

• Added --attention_backend CLI argument with choices: fa3, flashinfer, triton, fa4, flashmla, cutlass (default: fa3)
• Added --page_size CLI argument (default: 1)
• Added --mem_fraction_static CLI argument (default: 0.8)
• Removed the hardcoded module-level constants
• Updated start_server() to use args values
• Added attention backend info to the run-mode summary printout

Usage

# For RTX 4090 (Ada) users:
python infer.py --pdf document.pdf --attention_backend flashinfer
# Or with triton:
python infer.py --pdf document.pdf --attention_backend triton

Related

Fixes the minor secondary note in #46.

[kushdab]

Exactly the right fix -- the hardcoded fa3 was the last remaining "just edit the source" hurdle for non-Hopper users, and we flagged it explicitly in #46. Clean PR. A few notes:

Confirmed-valid backends for the custom SGLang wheel

The custom wheel (sglang-0.0.0.dev11416+g92e8bb79e) supports fa3, flashinfer, and triton. The other three in the choices list need verification:

• flashmla: Flash MLA (Multi-head Latent Attention) is actually the most architecturally relevant backend for this model -- UnlimitedOCR inherits DeepSeekV2's MLA mechanism, which MLA-aware backends exploit by not materializing the full KV cache. If the custom wheel includes it, flashmla should outperform fa3 on Hopper by reducing KV-cache memory footprint. Check with:

  python -c "from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend; print('available')"

• fa4: FlashAttention-4 (sm_90+). May not be compiled into this wheel version.
• cutlass: Available in some SGLang builds (from sglang.srt.layers.attention.cutlass_mla_backend import CutlassMLABackend).

argparse will happily accept any string in choices at the CLI level, but SGLang will crash at server startup if the backend isn't compiled in. A lightweight guard in start_server() prevents confusing failures:

_WHEEL_CONFIRMED = {"fa3", "flashinfer", "triton"}
if args.attention_backend not in _WHEEL_CONFIRMED:
    print(
        f"Warning: --attention_backend {args.attention_backend!r} is accepted by argparse "
        "but may not be compiled into the custom SGLang wheel. "
        "Server will fail at startup if unavailable.",
        file=sys.stderr,
    )

GPU-to-backend table in README

Users won't know which value to use without a guide. Suggested addition alongside the --attention_backend docs:

GPU                   SM      Recommended --attention_backend
---------------------------------------------------------------
H100 / H200           sm_90   fa3 (default) or flashmla (if wheel supports it)
RTX 4090 / L40S       sm_89   flashinfer
A100 / A800 / A30     sm_80   flashinfer or triton
RTX 3090 / A10        sm_86   flashinfer or triton
T4 / V100 / older     sm_75-  triton

Coordinate with PR #29

PR #29 (open) also modifies start_server() to add --trust-remote-code. Both PRs touch the same function -- either merge in order (merge #29 first, rebase #50 on top) or combine them to avoid a conflict at review time.

parse_args() refactor

Extracting argument parsing into a standalone parse_args() function was the right call -- it's exactly what makes the unit tests in tests/test_infer.py possible without launching a live server. Clean addition.