Optimize JitFunction direct CallState dispatch

[yhl-amd]

Summary

• avoid wrapping runtime arguments only to build the JIT cache key
• cache a direct CallState hot path for JitFunctions without self, Constexpr, or type parameters
• support kwargs/default reconstruction when using the direct call path

Context

This is the FlyDSL side of the DeepSeek-V4-Pro non-DPA 1k/1k c=256 regression fix. The regression was dominated by host-side launcher overhead in prefill attention-adjacent FlyDSL launchers, not by generated GPU kernel quality.

Paired with the AITER qk launcher patch, local ATOM E2E recovered most of the regression:

case	output tok/s	TPOT
old good	4250.48	58.04 ms
bad new	3907.11	63.32 ms
paired fix	4227.06	58.30 ms

Validation

• python -m py_compile python/flydsl/compiler/jit_function.py tests/unit/test_jit_direct_call_state.py
• PYTHONPATH=python python -m pytest tests/unit/test_jit_direct_call_state.py -q is blocked in my local source tree because it is not built and flydsl._mlir is absent; CI should run it after the normal build step.
• ATOM DeepSeek-V4-Pro non-DPA c=256 E2E with paired AITER patch: 4227.06 output tok/s, 58.30 ms TPOT.

[yhl-amd]

Added a FlyDSL-only hotpath replay benchmark in commit 610d9f94:

benchmarks/dsv4_hotpath_repro/
  bench_flydsl_hotpath.py
  run_matrix.sh
  README.md
  baselines.md

Purpose: reproduce the host-side FlyDSL launch overhead that showed up in the DeepSeek-V4-Pro non-DPA 1k/1k c=256 regression, without requiring AITER, ATOM, model weights, or a serving stack. The benchmark uses synthetic FlyDSL kernels with AITER-like launcher signatures and rank-local trace call counts:

qk_norm_rope_quant_like:      81 calls/window
fused_compress_attn_like:     80 calls/window
hca_compress_forward_like:    41 calls/window
hca_norm_rope_scatter_like:   41 calls/window

How to run a quick comparison from this PR checkout:

benchmarks/dsv4_hotpath_repro/run_matrix.sh \
  --original /path/to/original-flydsl \
  --fixed /path/to/fixed-flydsl \
  -- --windows 1 --warmup-windows 0 --gpu-event-calls 2

--fixed defaults to the current checkout, so after building this branch you can also run:

benchmarks/dsv4_hotpath_repro/run_matrix.sh \
  --original /path/to/original-flydsl \
  -- --windows 1 --warmup-windows 0 --gpu-event-calls 2

The helper detects these layouts for each path:

/path/to/FlyDSL/build-fly/python_packages
/path/to/FlyDSL/build/python_packages
/path/to/FlyDSL/python
/path/to/FlyDSL

For a longer replay closer to the original c=256 profile, drop the smoke args and use the defaults (--windows 16).

Local smoke result on gfx950 in the yhl_dev ROCm container:

mixed jit host time:
  original FlyDSL 0.2.0-pristine:      104.23 us/call
  fixed directraw-kw snapshot:           91.36 us/call
  speedup:                                1.14x

mixed compiled-position call time:
  original:                              10.09 us/call
  fixed:                                 10.01 us/call

qk_norm_rope_quant_like jit host time:
  original:                              68.68 us/call
  fixed:                                 64.89 us/call
  speedup:                                1.06x

attention-adjacent jit host time:
  fused_compress:                       152.45 -> 129.27 us/call  (1.18x)
  hca_compress:                         106.40 ->  90.80 us/call  (1.17x)
  hca_scatter:                           91.67 ->  78.31 us/call  (1.17x)

Interpretation: compiled_positional stays around ~10 us/call in both stacks, while the warm @flyc.jit path is much higher and improves with the hotpath fix. That isolates the regression to host-side JIT/cache-key/TensorAdaptor/CallState overhead, not GPU kernel math. This is intentionally not MoE-only; qk norm quant and the attention-adjacent launchers are included because those were the high-frequency calls in the original c=256 trace.